<jupyter_start><jupyter_text>Tutorial - Conditional LSTM Language Models>Training and using conditional LSTM language models  LSTM Language ModelsLSTM language models are a type of autoregressive generative model. This particular type of model is a good fit for RL-based optimization as they are light, robust and easy to optimize. These models make use of the [LSTM](https://arxiv.org/abs/1402.1128) architecture design.Language models are trained in a self-supervised fashion by next token prediction. Given a series of tokens, the model predicts a probability distribution over he next token. Self supervised training is very fast and doesn't require any data labels. Each text string in the dataset labels itself.During generation, we sample from the model in an autoregression fashion. Given an input token, the model predicts aa distribution of tokens over the next token. We then sample from that distributiona and feed the selected token back into the model. We repeat this process until either an end of sentence (EOS) token is predicted, or the generated sequence reaches a maximum allowed length.During sampling, we save the log probability of each token predicted. This gives us a probability value for the model's estimated likelihood of the generated compound. We can also backprop through this value. Conditional Language ModelsConditional language models condition the generated sequences on some latent vector. Conditioning can be implemented in two ways:- hidden conditioning: use the latent vector to initialize the hidden state of the model- output conditioning: concatenate the latent vector to the model activations right before the LSTM layersThe latent vector itself is generated by some encoder<jupyter_code>import sys
sys.path.append('..')

from mrl.imports import *
from mrl.core import *
from mrl.chem import *
from mrl.torch_imports import *
from mrl.torch_core import *
from mrl.layers import *
from mrl.dataloaders import *
from mrl.g_models.all import *
from mrl.train.agent import *
from mrl.vocab import *<jupyter_output>/home/dmai/miniconda3/envs/mrl/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: to-Python converter for boost::shared_ptr<RDKit::FilterCatalogEntry const> already registered; second conversion method ignored.
  return f(*args, **kwds)<jupyter_text>SetupBefore creating a model, we need to set up our data.Our raw data is in the form of SMILES strings. We need to convert these to tensors.First we need a `Vocab` to handle converting strings to tokens and mapping those tokens to integers. We will use the `CharacterVocab` class with the `SMILES_CHAR_VOCAB` vocabulary. This will tokenize SMILES on a character basis.More sophisticated tokenization schemes exist, but character tokenization is nice for the simplicity. Character tokenization has a small, compact vocabulary. Other tokenization strategies can tokenize by more meaningful subwords, but these strategies create a long tail of low frequency tokens and lots of `unk` characters.<jupyter_code># if in the repo
df = pd.read_csv('../files/smiles.csv')

# if in Collab:
# download_files()
# df = pd.read_csv('files/smiles.csv')
df.head()
vocab = CharacterVocab(SMILES_CHAR_VOCAB)<jupyter_output><empty_output><jupyter_text>`vocab` first tokenizes smiles into characters, then numericalizes the tokens into integer keys<jupyter_code>' '.join(vocab.tokenize(df.smiles.values[0]))
' '.join([str(i) for i in vocab.numericalize(vocab.tokenize(df.smiles.values[0]))])<jupyter_output><empty_output><jupyter_text>DatasetNow we need a dataset. For a conditional language model, we have to decide what data we want to send to the encoder to generate a latent vector. For this tutorial, we are going to use molecular fingerprints to generate the latent vector. The model will use the encoder to map the input fingerprint to a latent vector. Then the latent vector will be used to condition the hidden state of the decoder, which will reconstruct the SMILES string corresponding to the fingerprint.To do this, we will use the `Vec_To_Text_Dataset` dataset with `ECFP6` as our fingerprint function<jupyter_code>dataset = Vec_To_Text_Dataset(df.smiles.values, vocab, ECFP6)
dataloader = dataset.dataloader(32)<jupyter_output><empty_output><jupyter_text>Now we can look at the actual data<jupyter_code>x,y = next(iter(dataloader))
x
y<jupyter_output><empty_output><jupyter_text>The `x` tensor is a tuple containing `(fingerprint, smiles_ints)`. The fingerprint (`x[1]`) will be used to generate the latent vector, while the integer-coded SMILES string (`x[0]`) will be sent to the decoder along with the latent vector.You will notice the `y` tensor is the same as the `x[0]` tensor with the values shifted by one. This is because the goal of autoregressive language modeling is to predict the next character given the previous series of characters.  Model CreationWe can create a model through the `Conditional_LSTM_LM` class.First we need an encoder. Since our encoder data is a fingerprint, we will use a MLP-type encoder and set our latent vector to have a dimension of `512`<jupyter_code>enc_drops = [0.1, 0.1]
d_latent = 512

encoder = MLP_Encoder(2048, [1024, 512], d_latent, enc_drops)<jupyter_output><empty_output><jupyter_text>Now we create the model<jupyter_code>d_vocab = len(vocab.itos)
bos_idx = vocab.stoi['bos']
d_embedding = 256
d_hidden = 1024
n_layers = 3
tie_weights = True    
input_dropout = 0.3
lstm_dropout = 0.3
 
model = Conditional_LSTM_LM(
                            encoder,
                            d_vocab,
                            d_embedding,
                            d_hidden,
                            d_latent,
                            n_layers,
                            input_dropout=input_dropout,
                            lstm_dropout=lstm_dropout,
                            norm_latent=True,
                            condition_hidden=True,
                            condition_output=False,
                            bos_idx=bos_idx,
                        )<jupyter_output><empty_output><jupyter_text>We can examine the encoder-decoder structure of the model. We have a `MLP_Encoder` section which contains three MP layers. then we have a `decoder` section that consists of an embedding, three LSTM layers and an output layer<jupyter_code>model<jupyter_output><empty_output><jupyter_text>Now we'll put the model into a `GenerativeAgent` to manage supervised training.We need to specify a loss function - we will use standard cross entropy<jupyter_code>loss_function = CrossEntropy()
agent = GenerativeAgent(model, vocab, loss_function, dataset, base_model=False)<jupyter_output><empty_output><jupyter_text>Now we can train in a supervised fashion on next token prediction<jupyter_code>agent.train_supervised(32, 1, 1e-3)<jupyter_output><empty_output><jupyter_text>This was just a quick example to show the training API. We're not going to do a whole training process here. To train custom models, repeat this code with your own set of SMILES.  Pre-trained ModelsThe MRL model zoo offers a number of pre-trained models. We'll load one of these to continue.We'll use the `LSTM_LM_Small_Chembl` model. This model was trained first on a chunk of the ZINC database<jupyter_code>del model
del agent
gc.collect()
torch.cuda.empty_cache()
from mrl.model_zoo import FP_Cond_LSTM_LM_Small_ZINC
agent = FP_Cond_LSTM_LM_Small_ZINC(drop_scale=0.5, base_model=False)<jupyter_output><empty_output><jupyter_text>Now with a fully pre-trained model, we can look at drawing samples<jupyter_code>preds, lps = agent.model.sample_no_grad(256, 100, temperature=1.)
preds
lps.shape<jupyter_output><empty_output><jupyter_text>The `sample_no_grad` function gives is two outputs - `preds` and `lps`.`preds` is a long tensor of size `(bs, sl)` containing the integer tokens of the samples.`lps` is a float tensor of size `(bs, sl)` containing the log probabilities of each value in `preds`We can now reconstruct the predictions back into SMILES strings<jupyter_code>smiles = agent.reconstruct(preds)
smiles[:10]
mols = to_mols(smiles)<jupyter_output><empty_output><jupyter_text>Now lets look at some key generation statistics. - diversity - the percentage of unique samples- valid - the number of chemically valid samples<jupyter_code>div = len(set(smiles))/len(smiles)
val = len([i for i in mols if i is not None])/len(mols)
print(f'Diversity:\t{div:.3f}\nValid:\t\t{val:.3f}')
valid_mols = [i for i in mols if i is not None]
draw_mols(valid_mols[:16], mols_per_row=4)<jupyter_output><empty_output><jupyter_text>Conditional GenerationLets look further at how conditional generation works. One important aspect of conditioning on a latent space is the ability to sample from it. This poses an interesting challenge for a conditional language model, because unlike models like `VAE`, there is no set prior distribution. The `Conditional_LSTM_LM` class handles this with the `norm_latent` argument. If `norm_latent=true` (as above), the latent vectors are all normalized to a magnitude of 1 before being sent to the decoder. This is similar to what is done in StyleGan-type models.By setting the norm constraint, we can sample valid latent vectors by sampling from a normal distribution and normalizing the resulting vectors.<jupyter_code>latents = to_device(torch.randn((64, agent.model.encoder.d_latent)))
preds, lps = agent.model.sample_no_grad(latents.shape[0], 100, z=latents, temperature=1.)
smiles = agent.reconstruct(preds)
mols = to_mols(smiles)
mols = [i for i in mols if i is not None]
draw_mols(mols[:16], mols_per_row=4)<jupyter_output><empty_output><jupyter_text>Conditional Samples on a Single InputOne advantage of using a conditional generation is the ability to sample multiple compounds from a specific input.Here we get a latent vector from a specific molecule in the datase<jupyter_code>smile = df.smiles.values[0]
to_mol(smile)
smile_ds = dataset.new([smile])
batch = to_device(smile_ds.collate_function([smile_ds[0]]))
x,y = batch
agent.model.eval();
latents = agent.model.x_to_latent(x)
latents.shape<jupyter_output><empty_output><jupyter_text>Now we sample 64 molecules from the same latent vector. The diversity of the outputs will depend on the sampling temperature used and the degree of dropout (if enabled)<jupyter_code>agent.model.train();
preds, lps = agent.model.sample_no_grad(64, 100, z=latents.repeat(64,1), temperature=1.)
smiles = agent.reconstruct(preds)
print(len(set(smiles)), len(set(smiles))/len(smiles))
smiles = list(set(smiles))
mols = to_mols(smiles)
mols = [i for i in mols if i is not None]
len(mols)
draw_mols(mols[:16], mols_per_row=4)<jupyter_output><empty_output><jupyter_text>In the above section, we generated 64 compounds and ended up with 10 unique SMILES strings.If we want more diversity, we can sample with a higher temperature.You'll see we get 40 unique SMILES. However, only 36 of them are valid SMILES strings. This is the tradeoff of using higher sampling temperatures<jupyter_code>preds, lps = agent.model.sample_no_grad(64, 100, z=latents.repeat(64,1), temperature=2.)
smiles = agent.reconstruct(preds)
print(len(set(smiles)), len(set(smiles))/len(smiles))
smiles = list(set(smiles))
mols = to_mols(smiles)
mols = [i for i in mols if i is not None]
len(mols)
draw_mols(mols[:16], mols_per_row=4)<jupyter_output><empty_output><jupyter_text>PriorThe `Conditional_LSTM_LM` supports a defined prior distrribution for sampling. By default, this prior is not enabled.We can set the prior using the latent vector we generated earlier to focus the model's prior distribution on that area of the latent space.To create this distribution, we first need to define the log-variance of the latent space. We will set the log variance to `-5` to set aa tight distribution around the latent vector.<jupyter_code>logvar = to_device(torch.zeros(latents.shape)-5)
agent.model.set_prior_from_latent(latents.squeeze(0), logvar.squeeze(0), trainable=False)
preds, lps = agent.model.sample_no_grad(128, 100, temperature=1.)
smiles = agent.reconstruct(preds)
print(len(set(smiles)), len(set(smiles))/len(smiles))
smiles = list(set(smiles))
mols = to_mols(smiles)
mols = [i for i in mols if i is not None]
len(mols)
draw_mols(mols[:16], mols_per_row=4)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.0761144161030071, 'fraction_numerical': 0.011246027218645586, 'mean_word_length': 4.9921875, 'pattern_counts': {'":': 0, '<': 77, '<?xml version=': 0, '>': 78, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '19', 'dir': 'jupyter-structured-clean-dedup', 'id': '12178', 'license': 'MIT', 'n_tokens_mistral': 3822, 'n_tokens_neox': 3606, 'n_words': 1390, 'path': 'nbs/tutorials/tutorials.generative_models.conditional_lstm_lm.ipynb', 'repo_name': 'DarkMatterAI/mrl', 'size': '12271'}