Adding simplified data wrangling files

Author: szaman19

applications/FLASK/Transformer/README.md

@@ -0,0 +1,11 @@
+# Molecular Transformer Baseline
+
+This is an encoder-decoder transformer on SMILES data
+
+## Preparing the data
+
+The data can be prepared ahead of time to simplify the pipeline
+
+---
+This example is a simplified encoder-decoder transformer in the `application/nlp` directory. 
+

applications/FLASK/Transformer/arg_utils.py

@@ -0,0 +1,52 @@
+import argparse
+import data_utils
+
+
+# Command-line arguments
+def add_transformer_architecture_arguments(args: argparse.Namespace):
+    """
+    Adds the command line arguments to specify transformer architecture model
+    parameters. This is only relevant for the encoder-decoder transformer model.
+    """
+    args.add_argument('--num-attention-heads',
+                      action='store',
+                      default=8,
+                      type=int,
+                      help='number of parallel attention layers (default: 8)',
+                      metavar='NUM')
+    args.add_argument('--embed-dim',
+                      action='store',
+                      default=512,
+                      type=int,
+                      help='embedding space dimension (default: 512)',
+                      metavar='NUM')
+    args.add_argument('--feedforward-dim',
+                      action='store',
+                      default=0,
+                      type=int,
+                      help='feedforward network dimension. If zero, set to be '
+                      '4 times the embedding dimension (default: 0)',
+                      metavar='NUM')
+    args.add_argument('--num-layers',
+                      action='store',
+                      default=6,
+                      type=int,
+                      help='Number of encoder and decoder layers (default: 6)',
+                      metavar='NUM')
+
+
+def add_dataset_arguments(args: argparse.Namespace, default: str):
+    """
+    Adds dataset-related arguments to an existing argparse object.
+    """
+    args.add_argument('--dataset',
+                      type=str,
+                      default=default,
+                      help=f'Which dataset to use (default: {default})',
+                      choices=data_utils.available_datasets())
+    args.add_argument('--dataset-fraction',
+                      action='store',
+                      default=1.0,
+                      type=float,
+                      help='Fraction of dataset to use (default: 1.0)',
+                      metavar='NUM')

applications/FLASK/Transformer/dataset_utils.py

@@ -0,0 +1,25 @@
+import importlib
+import os
+import sys
+from typing import List
+
+dataset_dir = os.path.join(os.path.dirname(__file__), 'datasets')
+
+
+def available_datasets() -> List[str]:
+    """
+    Returns the available datasets in the dataset folder.
+    """
+    result = []
+    for file in os.listdir(dataset_dir):
+        if file.endswith('.py'):
+            result.append(os.path.basename(file)[:-3])
+    return result
+
+
+def load_dataset(name: str):
+    """
+    Loads a dataset by importing the requested module.
+    """
+    sys.path.append(dataset_dir)
+    return importlib.import_module(name)

applications/FLASK/Transformer/datasets/QM9.py

@@ -0,0 +1,55 @@
+"""
+The QM9 dataset, stored as pre-tokenized binary files for optimized processing.
+"""
+import os
+import os.path
+import pickle
+
+import numpy as np
+from pretokenize.SMILES_tokenizer import MolTokenizer
+
+sequence_length = int(os.getenv('QM9_SEQUENCE_LENGTH', default='32'))
+
+# ----------------------------------------------
+# Setup
+# ----------------------------------------------
+
+# Load the datasets
+data_dir = os.getenv(
+    'QM9_DATA_DIR',
+    '/p/vast1/lbann/datasets/FLASK/qm9')
+
+tokenizer = MolTokenizer("SMILES_vocab.json")
+tokenizer.load_vocab_file()
+
+dataset_train = np.load(os.path.join(data_dir, 'QM9_Pretokenize.py'))
+
+_vocab_size = 46
+
+pad_index = tokenizer.token_to_id('<pad>')
+bos_index = tokenizer.token_to_id('<bos>')
+eos_index = tokenizer.token_to_id('<eos>')
+
+# ----------------------------------------------
+# Sample access functions
+# ----------------------------------------------
+
+def num_train_samples():
+    return dataset_train.shape[0]
+
+def get_train_sample(i):
+    data = dataset_train[i]
+
+    return 
+
+def sample_dims():
+    return (2 * sequence_length + 1, )
+
+def vocab_size():
+    return _vocab_size
+
+
+if __name__ == '__main__':
+    print('Training samples:', num_train_samples())
+    print('Training sample 101:')
+    print(get_train_sample(101))

applications/FLASK/Transformer/datasets/pretokenize/QM9_Pretokenize.py

@@ -0,0 +1,30 @@
+import numpy as np
+from SMILES_tokenizer import MolTokenizer
+
+
+def random_zero_array(arr, probability, mask):
+    return np.where(np.random.random(arr.shape) < probability, mask, arr)
+
+
+def main():
+    tokenizer = MolTokenizer("SMILES_vocab.json")
+    tokenizer.load_vocab_file()
+    with open("QM9_smiles.txt", 'r') as smiles_data:
+      smiles_data = smiles_data.readlines()
+      num_samples = len(smiles_data)
+      max_length = 32
+
+      tokenized_data = np.ones((num_samples, max_length)) * tokenizer.encode(tokenizer.pad_token) 
+      tokenized_data[:, 0] = tokenizer.encode(tokenizer.sep_token)
+
+      for i, smiles in enumerate(smiles_data, start=1):
+        tokens = tokenizer.tokenize(smiles)
+        tokens = random_zero_array(tokens, 0.15, tokenizer.encode(tokenizer.mask_token))
+        tokenized_data[i, :len(tokens)] = tokens
+        tokenized_data[i, len(tokens)] = tokenizer.encode(tokenizer.cls_token)
+
+    np.save('QM9_Pretokenized.npy', tokenized_data)
+
+if __name__ == '__main__':
+    main()
+