Upload finetune_gpt.py

finetune_gpt.py

@@ -0,0 +1,442 @@
+import deepchem as dc
+import tensorflow as tf
+import numpy as np
+import random
+import pandas as pd
+from rdkit import Chem
+from rdkit.Chem import Draw
+import os
+
+def finetune_gpt(df, chembl_id):
+  '''
+  accepts a dataframe with SMILES and uses deepchem to tokenize the dataset, 
+  then uses tensorflow and a pre-trained model to fine tune the model on the dataset.
+  The pretrained model was trained on 305K molecules from the ZN15 dataset, including at least
+  50K that are bioactive. 
+
+  Returns:
+    out_text: the generated molecules
+    img: the image of the generated molecules
+
+  requires files:
+    vocab.txt
+    vocab_305K.txt
+    GPT_ZN305_50epochs.weights.h5
+    layer_store_GPT_ZN305_50epochs.txt
+    ZN305K_smiles.csv
+
+  '''
+  # chemck to see if f"gen_smiles_{chembl_id}.csv" exists
+  if os.path.exists(f"gen_smiles_{chembl_id}.csv"):
+    df = pd.read_csv(f"gen_smiles_{chembl_id}.csv")
+    final_smiles = df["SMILES"].to_list()
+    final_mols = [Chem.MolFromSmiles(smile) for smile in final_smiles]
+  else:
+
+    # Prepare dataset from chembl ==========================================
+
+    if len(df) > 2000:
+      df = df.sample(n=2000, random_state=42)
+
+    smiles_list = df["SMILES"].to_list()
+
+    Xa = []
+    for smiles in smiles_list:
+      smiles = smiles.replace("[Na+].","").replace("[Cl-].","").replace(".[Cl-]","").replace(".[Na+]","")
+      smiles = smiles.replace("[K+].","").replace("[Br-].","").replace(".[K+]","").replace(".[Br-]","")
+      smiles = smiles.replace("[I-].","").replace(".[I-]","").replace("[Ca2+].","").replace(".[Ca2+]","")
+      Xa.append(smiles)
+
+    tokenizer=dc.feat.SmilesTokenizer(vocab_file="vocab.txt")
+    featname="SMILES Tokenizer"
+
+    fl = list(map(lambda x: tokenizer.encode(x),Xa))
+
+    biggest = 1
+    smallest = 200
+    for i in range(len(fl)):
+        temp = len(fl[i])
+        if temp > biggest:
+            biggest = temp
+        if temp < smallest:
+            smallest = temp
+
+    print(biggest, smallest)
+
+    string_length = smallest - 1
+    max_length = biggest
+
+    fl2 = list(map(lambda x: tokenizer.add_padding_tokens(x,max_length),fl))
+
+    fl2set=set()
+    for sublist in fl2:
+      fl2set.update(sublist)
+    new_vocab_size = len(fl2set)
+    print("New vocabulary size: ",new_vocab_size)
+
+    f = open("vocab_305K.txt", "r")
+    raw_lines = f.readlines()
+    f.close()
+    VOCAB_SIZE = len(raw_lines)
+    print("Vocabulary size for standard dataset: ",VOCAB_SIZE)
+
+    lines = []
+    for line in raw_lines:
+      lines.append(line.replace("\n",""))
+
+    novel_items = []
+    for item in fl2set:
+      item = tokenizer.decode([item])
+      item = tokenizer.convert_tokens_to_string(item)
+      item = item.replace(" ","")
+
+      if item not in lines:
+        print(f"{item} not in standard vocabulary")
+        novel_items.append(item)
+
+    if(len(novel_items) > 0):
+      print("This dataset is not compatible with the Foundation model vocabulary")
+    else:
+      print("This dataset is compatible with the Foundation model vocabulary")
+
+    if max_length > 166:
+      print("This dataset's context window is not compatible with the Foundation model.")
+    else:
+      print("This dataset's context window is compatible with the Foundation model")
+
+    smiles_removed_tokens = []
+    for i,smiles in enumerate(Xa):
+      bad_list = [True if (token in smiles) else False for token in novel_items]
+      if not any(bad_list):
+        smiles_removed_tokens.append(smiles)
+
+    smiles_no_long = []
+    for i,smiles in enumerate(smiles_removed_tokens):
+      if len(smiles) <= 166:
+        smiles_no_long.append(smiles)
+
+    print(f"Removed {len(Xa) - len(smiles_no_long)} entries from the list!")
+
+    new_dict = {"SMILES": smiles_no_long}
+    new_df = pd.DataFrame(new_dict)
+
+    Xa = []
+    for smiles in new_df['SMILES']:
+      Xa.append(smiles)
+
+    tokenizer=dc.feat.SmilesTokenizer(vocab_file="vocab_305K.txt")
+    featname="SMILES Tokenizer"
+
+    fl = list(map(lambda x: tokenizer.encode(x),Xa))
+
+    biggest = 1
+    smallest = 200
+    for i in range(len(fl)):
+        temp = len(fl[i])
+        if temp > biggest:
+            biggest = temp
+        if temp < smallest:
+            smallest = temp
+
+    print(biggest, smallest)
+
+    string_length = smallest - 1
+    max_length = biggest
+
+    fl2 = list(map(lambda x: tokenizer.add_padding_tokens(x,max_length),fl))
+
+    f = open("vocab_305K.txt", "r")
+    lines = f.readlines()
+    f.close()
+    VOCAB_SIZE = len(lines)
+    print("Vocabulary size for this dataset: ",VOCAB_SIZE)
+
+    x = []
+    y = []
+    i=0
+    for string in fl2:
+        x.append(string[0:max_length-1]) #string_length
+        y.append(string[1:max_length]) #string_length+1
+
+    fx = np.array(x)
+    fy = np.array(y)
+    print("Number of features and datapoints, targets: ",fx.shape,fy.shape)
+
+    # Load foundation model ==================================================
+
+    VOCAB_SIZE = 100
+    max_length = 166
+    num_new_blocks = 2
+    EMBEDDING_DIM = 256
+    N_HEADS = 4
+    KEY_DIM = 256
+    FEED_FORWARD_DIM = 256
+
+    inputs = tf.keras.layers.Input(shape=(None,),dtype=tf.int32)
+    x = TokenAndPositionEmbedding(max_length,VOCAB_SIZE,EMBEDDING_DIM)(inputs)
+    for i in range(num_new_blocks+2):
+      x, attentions_scores = TransformerBlock(N_HEADS,KEY_DIM,EMBEDDING_DIM,FEED_FORWARD_DIM)(x)
+    outputs = tf.keras.layers.Dense(VOCAB_SIZE,activation="softmax")(x)
+
+    gpt_ft = tf.keras.models.Model(inputs = inputs, outputs =[outputs, attentions_scores])
+
+    f = open("layer_store_GPT_ZN305_50epochs.txt", "r")
+    layer_name_store_raw = f.readlines()
+    f.close()
+
+    print("Reading in layers:")
+    layer_name_store = []
+    for line in layer_name_store_raw:
+        line = line.replace("\n","")
+        layer_name_store.append(line)
+        print(line)
+    print("===========================================")
+
+    new_layers = num_new_blocks + 1
+    for i,layer in enumerate(gpt_ft.layers[:-new_layers]):
+      layer.name = layer_name_store[i]
+      print(f"{layer.name} has been named!")
+
+    for i,layer in enumerate(gpt_ft.layers[-new_layers:-1]):
+      layer.name = f"transformer_block_X_{i+1}"
+      print(f"{layer.name} has been named!")
+
+    gpt_ft.layers[-1].name = "dense_X"
+
+    gpt_ft.load_weights("GPT_ZN305_50epochs.weights.h5", skip_mismatch=True)
+
+    for layer in gpt_ft.layers[0:-new_layers]:                 #make old layers freeze and only train new layers
+      layer.trainable=False
+      print(f"setting layer {layer.name} untrainable.")
+
+    for layer in gpt_ft.layers[-new_layers:]:
+      layer.trainable=True
+      print(f"setting layer {layer.name} trainable.")
+
+    # train new layers =======================================================
+
+    batch_size = 512
+    gpt_ft.compile("adam",loss=[tf.keras.losses.SparseCategoricalCrossentropy(),None])
+    gpt_ft.fit(fx,fy,epochs = 50, batch_size = batch_size)
+
+    # train all together =====================================================
+    for layer in gpt_ft.layers:
+      layer.trainable=True
+      print(f"setting layer {layer.name} trainable.")
+
+    gpt_ft.compile("adam",loss=[tf.keras.losses.SparseCategoricalCrossentropy(),None])
+    gpt_ft.fit(fx,fy,epochs = 25, batch_size = batch_size)
+
+    # make prompts ============================================================
+
+    df_prompts = pd.read_csv("ZN305K_smiles.csv")
+
+    Xap = []
+    for smiles in df_prompts["SMILES"]:
+      smiles = smiles.replace("[Na+].","").replace("[Cl-].","").replace(".[Cl-]","").replace(".[Na+]","")
+      smiles = smiles.replace("[K+].","").replace("[Br-].","").replace(".[K+]","").replace(".[Br-]","")
+      smiles = smiles.replace("[I-].","").replace(".[I-]","").replace("[Ca2+].","").replace(".[Ca2+]","")
+      Xap.append(smiles)
+
+    raw_prompts = random.choices(Xap,k=50)
+
+    test_string = []
+    for smile in raw_prompts:
+      test_string.append(smile[:2])
+
+    # inference ================================================================
+
+    tf.random.set_seed(42)
+
+    batch_length = len(test_string)
+    prompt_length = len(test_string[0])
+    test_xlist = np.empty([batch_length,prompt_length], dtype=int)
+
+    test_tokenized = list(map(lambda x: tokenizer.encode(x),test_string))
+    for i in range(batch_length):
+        test_xlist[i][:] = test_tokenized[i][:prompt_length]
+    test_array = np.array(test_xlist)
+
+    proba = np.empty([batch_length,VOCAB_SIZE])
+    rescaled_logits = np.empty([batch_length,VOCAB_SIZE])
+    preds = np.empty([batch_length])
+    gen_molecules = np.empty([batch_length])
+
+    c_final = 60 - prompt_length
+    sig_start = 0.10
+    TEMP = 1.5
+
+    for c in range(0,c_final,1):
+
+        c_o = int(c_final*sig_start)
+
+        T_int = TEMP*(1/(1+np.exp(-(c-c_o))))
+
+        results, _ = gpt_ft.predict(test_array)
+
+        if T_int < 0.015:
+            print(f"using zero temp generation with {T_int}.")
+            for j in range(batch_length):
+                preds[j] = tf.argmax(results[j][-1])
+                preds = list(map(lambda x: int(x),preds))
+        else:
+            print(f"using variable temp generation with {T_int}.")
+            for j in range(batch_length):
+                proba[j] = (results[j][-1:]) ** (1/T_int)
+                rescaled_logits[j] = ( proba[j][:] ) / np.sum(proba[j][:])
+                preds[j] = np.random.choice(len(rescaled_logits[j][:]),
+                                            p=rescaled_logits[j][:])
+                preds = list(map(lambda x: int(x),preds))
+        test_array = np.c_[test_array,preds]
+        print(test_array.shape)
+
+    gen_molecules = list(map(lambda x: tokenizer.decode(x),test_array))
+    gen_molecules = list(map(lambda x: tokenizer.convert_tokens_to_string(x),
+                              gen_molecules))
+    gen_molecules = list(map(lambda x: strip_smiles(x),gen_molecules))
+
+    mols, smiles = mols_from_smiles(gen_molecules)
+
+    final_smiles = []
+    final_mols = []
+    for smile, mol in zip(smiles,mols):
+        if smile not in final_smiles:
+            final_smiles.append(smile)
+            final_mols.append(mol)
+    
+    final_dict = {"SMILES": final_smiles}
+    final_df = pd.DataFrame.from_dict(final_dict)
+    final_df.to_csv(f"gen_smiles_{chembl_id}.csv", index = False)
+
+  print(f"Generated {len(final_smiles)} unique molecules.")
+
+  img = Draw.MolsToGridImage(final_mols,molsPerRow=3,legends=final_smiles)
+  #img.save("Substitution_image.png")
+
+  out_text = f'The generated molecules are: \n'
+  for smile in final_smiles:
+    out_text += f'{smile}\n'
+
+  return out_text, img
+
+def casual_attention_mask(batch_size,n_dest,n_src,dtype):
+  '''
+    Make a causal attention mask
+  '''
+  i = tf.range(n_dest)[:,None]
+  j = tf.range(n_src)
+  m = i >= j - n_src + n_dest
+  mask = tf.cast(m,dtype)
+  mask = tf.reshape(mask,[1,n_dest,n_src])
+  mult = tf.concat([tf.expand_dims(batch_size,-1),tf.constant([1,1],dtype=tf.int32)],0)
+  return tf.tile(mask,mult)
+
+class TransformerBlock(tf.keras.layers.Layer):
+  '''
+    Transformer block with multi-head attention.
+  '''
+  def __init__(self,num_heads,key_dim,embed_dim,ff_dim,dropout_rate=0.1):
+    super(TransformerBlock,self).__init__()
+    self.num_heads = num_heads
+    self.key_dim = key_dim
+    self.embed_dim = embed_dim
+    self.ff_dim = ff_dim
+    self.dropout_rate = dropout_rate
+    self.attn = tf.keras.layers.MultiHeadAttention(self.num_heads,self.key_dim,
+                                                    output_shape=self.embed_dim)
+    self.dropout_1 = tf.keras.layers.Dropout(self.dropout_rate)
+    self.ln_1 = tf.keras.layers.LayerNormalization(epsilon=0.000001)
+    self.ffn_1 = tf.keras.layers.Dense(self.ff_dim,activation="relu")
+    self.ffn_2 = tf.keras.layers.Dense(self.embed_dim)
+    self.dropout_2 = tf.keras.layers.Dropout(self.dropout_rate)
+    self.ln_2 = tf.keras.layers.LayerNormalization(epsilon=0.000001)
+
+  def call(self,inputs):
+    input_shape = tf.shape(inputs)
+    batch_size2 = input_shape[0]
+    seq_len = input_shape[1]
+    casual_mask = casual_attention_mask(batch_size2,seq_len,seq_len,tf.bool)
+    attention_output, attention_scores = self.attn(inputs,inputs,
+                                                    attention_mask=casual_mask,
+                                                    return_attention_scores=True)
+    attention_output = self.dropout_1(attention_output)
+    out1 = self.ln_1(inputs + attention_output)
+    ffn_1 = self.ffn_1(out1)
+    ffn_2 = self.ffn_2(ffn_1)
+    ffn_output = self.dropout_2(ffn_2)
+    return (self.ln_2(out1+ffn_output),attention_scores)
+
+  def get_config(self):
+    config = super().get_config()
+    config.update({"key_dim": self.key_dim, "embed_dim": self.embed_dim,
+                  "num_heads": self.num_heads,"ff_dim": self.ff_dim,
+                  "dropout_rate": self.dropout_rate})
+    return config
+
+class TokenAndPositionEmbedding(tf.keras.layers.Layer):
+  '''
+    Embeds tokens and positions.
+  '''
+  def __init__(self,max_len,vocab_size,embed_dim):
+    super(TokenAndPositionEmbedding,self).__init__()
+    self.max_len = max_len
+    self.vocab_size = vocab_size
+    self.embed_dim = embed_dim
+    self.token_emb = tf.keras.layers.Embedding(input_dim=vocab_size,
+                                                output_dim = embed_dim)
+    self.pos_emb = tf.keras.layers.Embedding(input_dim=max_len,output_dim=embed_dim)
+
+  def call(self,x):
+    maxlen = tf.shape(x)[-1]
+    positions = tf.range(start=0,limit=maxlen,delta=1)
+    positions = self.pos_emb(positions)
+    x = self.token_emb(x)
+    return x + positions
+
+  def get_config(self):
+    config = super().get_config()
+    config.update({"max_len": self.max_len, "vocab_size": self.vocab_size,
+                  "embed_dim": self.embed_dim})
+    return config
+
+def strip_smiles(input_string):
+  '''
+    Cleans un-needed tokens from the SMILES string.
+
+      Args:
+        input_string: SMILES string
+      Returns:
+        output_string: cleaned SMILES string
+  '''
+  output_string = input_string.replace(" ","").replace("[CLS]","").replace("[SEP]","").replace("[PAD]","")
+  output_string = output_string.replace("[Na+].","").replace(".[Na+]","")
+  return output_string
+
+def mols_from_smiles(input_smiles_list):
+  '''
+    Converts a list of SMILES strings to a list of RDKit molecules.
+
+      Args:
+        input_smiles_list: list of SMILES strings
+      Returns:
+        valid_mols: list of RDKit molecules
+        valid_smiles: list of SMILES strings
+  '''
+  valid_mols = []
+  valid_smiles = []
+
+  good_count = 0
+  for ti, smile in enumerate(input_smiles_list):
+    temp_mol = Chem.MolFromSmiles(smile)
+    if temp_mol != None:
+      valid_mols.append(temp_mol)
+      valid_smiles.append(smile)
+      good_count += 1
+    else:
+      print(f"SMILES {ti} was not valid!")
+
+  if len(valid_mols) == len(valid_smiles) == good_count:
+    print(f"Generated a total of {good_count} mol objects")
+  else:
+    print("mismatch!")
+  return valid_mols, valid_smiles