Upload 2 files

Scripts for data prep and model training

prepare_data_script.py

@@ -0,0 +1,91 @@
+import pandas as pd
+import os
+
+import torch
+from transformers import RobertaTokenizerFast, RobertaForMaskedLM, DataCollatorWithPadding
+
+import datasets
+from datasets import disable_caching
+disable_caching()
+
+DEVICE = 'cuda:0'                                   # model device
+ENCODER_MODEL_NAME = "entropy/roberta_zinc_480m"    # encoder name
+ENCODER_BATCH_SIZE = 1024                           # batch size for computing embeddings
+
+TOKENIZER_MAX_LEN = 256                             # max_length param on tokenizer
+TOKENIZATION_NUM_PROC = 32                          # number of processes for tokenization
+
+'''
+Data source is expected to be a CSV file with a column of SMILES strings 
+denoted by `SMILES_COLUMN`. The CSV is processed in chunks of size `PROCESS_CHUNKSIZE`. 
+
+Processed chunks are saved to `SAVE_PATH` with the format `SAVE_PATH/processed_shard_{i}.hf`
+'''
+
+DATASET_CSV_FILENAME = None                         # path to data csv
+PROCESS_CHUNKSIZE = 1000000                         # how many rows to process/save for each dataset shard
+SMILES_COLUMN = 'smiles'                            # csv column holding smiles strings
+MAX_CHUNKS = None                                   # total number of chunks to process (if None, all chunks are processed)
+MAX_SMILES_LENGTH = 90                              # max smiles string length (exclusive)
+MIN_SMILES_LENGTH = 5                               # min smiles string length (exclusive)
+FILTER_NUM_PROC = 32                                # number of processes for filtering
+SAVE_PATH = None                                    # directory to save data shards to
+
+assert DATASET_CSV_FILENAME is not None, "must specify dataset filename"
+assert SAVE_PATH is not None, "must specify save path"
+
+
+def tokenization(example):
+    return tokenizer(example[SMILES_COLUMN], add_special_tokens=True, 
+                     truncation=True, max_length=TOKENIZER_MAX_LEN)
+
+def embed(inputs):
+    inputs = {k:inputs[k] for k in ['input_ids', 'attention_mask']}
+    inputs = collator(inputs)
+    inputs = {k:v.to(DEVICE) for k,v in inputs.items()}
+    
+    with torch.no_grad():
+        outputs = model(**inputs, output_hidden_states=True)
+        full_embeddings = outputs[-1][-1]
+        mask = inputs['attention_mask']
+        
+        mean_embeddings = ((full_embeddings * mask.unsqueeze(-1)).sum(1) / mask.sum(-1).unsqueeze(-1))
+        
+    return {'encoder_hidden_states' : mean_embeddings}
+
+def length_filter_smiles(example):
+    min_check = (len(example[SMILES_COLUMN])>MIN_SMILES_LENGTH) if (MIN_SMILES_LENGTH is not None) else True
+    max_check = (len(example[SMILES_COLUMN])<MAX_SMILES_LENGTH) if (MIN_SMILES_LENGTH is not None) else True
+    type_check = type(example[SMILES_COLUMN])==str
+    filter_pass = all([min_check, max_check, type_check])
+    return filter_pass
+
+
+tokenizer = RobertaTokenizerFast.from_pretrained(ENCODER_MODEL_NAME, max_len=TOKENIZER_MAX_LEN)
+collator = DataCollatorWithPadding(tokenizer, padding=True, return_tensors='pt')
+
+model = RobertaForMaskedLM.from_pretrained(ENCODER_MODEL_NAME)
+model.to(DEVICE)
+model.eval()
+
+df_iter = pd.read_csv(DATASET_CSV_FILENAME, chunksize=PROCESS_CHUNKSIZE, usecols=[SMILES_COLUMN])
+
+for i, df in enumerate(df_iter):
+    print(f'processing dataset chunk {i}')
+    
+    dataset = datasets.Dataset.from_pandas(df)
+    
+    dataset = dataset.filter(lambda example: length_filter_smiles(example), num_proc=FILTER_NUM_PROC)
+    
+    dataset = dataset.map(tokenization, batched=True, num_proc=TOKENIZATION_NUM_PROC)
+    
+    dataset = dataset.map(embed, batched=True, batch_size=ENCODER_BATCH_SIZE)
+    
+    dataset.save_to_disk(f'{SAVE_PATH}/processed_shard_{i}.hf')
+    
+    if (MAX_CHUNKS is not None) and (i >= MAX_CHUNKS-1):
+        break
+
+print('finished data processing')
+
+

train_script.py

@@ -0,0 +1,102 @@
+import pandas as pd
+import os
+
+import torch
+import torch.nn as nn
+
+from transformers import GPT2TokenizerFast, GPT2LMHeadModel
+from transformers import DataCollatorWithPadding, GPT2Config, DataCollatorForLanguageModeling
+from transformers import Trainer, TrainingArguments, RobertaTokenizerFast
+
+import datasets
+from datasets import disable_caching
+disable_caching()
+from datasets import IterableDataset
+
+from conditional_gpt2_model import ConditionalGPT2LMHeadModel
+
+
+ENCODER_MODEL_NAME = "entropy/roberta_zinc_480m"  # encoder model name
+TOKENIZER_MAX_LEN = 256                           # max_length param on tokenizer
+
+DATA_SUBSHARDS = 10                               # number of shards to break each data chunk into
+
+DATA_DIR = None                                   # directory with saved data shards
+TRAINER_SAVE_DIR = None                           # directory to save model checkpoints
+
+assert DATA_DIR is not None, "data directory must be specified"
+assert TRAINER_SAVE_DIR is not None, "trainer save directory must be specified"
+
+
+
+def gen_dataset():
+    
+    data_filenames = sorted([i for i in os.listdir(DATA_DIR) if '.hf' in i])
+    
+    for filename in data_filenames:
+        
+        dataset = datasets.Dataset.load_from_disk(f'{DATA_DIR}/{filename}')
+        
+        keep_cols = ['input_ids', 'encoder_hidden_states']
+        
+        dataset = dataset.remove_columns([i for i in dataset.column_names 
+                                          if not i in keep_cols]).with_format("torch")
+        
+        # contiguous shards for faster loading
+        shards = [dataset.shard(num_shards=DATA_SUBSHARDS, index=index, contiguous=True) 
+                  for index in range(DATA_SUBSHARDS)]
+        
+        for i, shard in enumerate(shards):
+            for example in shard:
+                # need to add unit axis to hidden states
+                example['encoder_hidden_states'] = example['encoder_hidden_states'][None,:]
+                yield example
+
+dataset = IterableDataset.from_generator(gen_dataset)
+dataset = dataset.with_format("torch")
+
+tokenizer = RobertaTokenizerFast.from_pretrained(ENCODER_MODEL_NAME, max_len=TOKENIZER_MAX_LEN)
+collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
+
+
+config = GPT2Config(
+    vocab_size=len(tokenizer),
+    n_positions=TOKENIZER_MAX_LEN,
+    bos_token_id=tokenizer.bos_token_id,
+    eos_token_id=tokenizer.eos_token_id,
+    n_layer=6,
+    n_head=8,
+    add_cross_attention=True,
+)
+
+model = ConditionalGPT2LMHeadModel(config)
+
+# change trainer args as needed
+args = TrainingArguments(
+    output_dir=TRAINER_SAVE_DIR,
+    per_device_train_batch_size=192,
+    logging_steps=25,
+    gradient_accumulation_steps=8,
+    num_train_epochs=1,
+    weight_decay=0.1,
+    warmup_steps=1000,
+    lr_scheduler_type="cosine",
+    learning_rate=1e-5,
+    save_steps=200,
+    save_total_limit=30,
+    fp16=True,
+    push_to_hub=False,
+    max_steps=50000,
+)
+
+
+trainer = Trainer(
+    model=model,
+    tokenizer=tokenizer,
+    args=args,
+    data_collator=collator,
+    train_dataset=dataset,
+)
+
+trainer.train()
+