{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Attention is What You Get\n\nThis is my entry into the [Bristol-Myers Squibb Molecular Translation](https://www.kaggle.com/c/bms-molecular-translation)  Kaggle competition.\n\n-----\n\nAUTHOR: \n\nMo Venouziou\n\n- *Email: mvenouziou@gmail.com*\n- *LinkedIn: www.linkedin.com/in/movenouziou/*\n\nUpdates:\n\n - *Original Posting: June 2, 2021*\n - *07/04/21: fixed kaggle compat issues*\n - *06/25/21: fixed optional Conv decoder layers*\n - *06/21/21: added TPU support*\n - *06/17/21: improved training & inference speed. Allows full AIAYN model size on TPU, faster small model training on GPU. (note: optional Conv decoder layers code broken)*\n\n----","metadata":{}},{"cell_type":"markdown","source":"### Our Goal: Predict the \"InChI\" value of any given chemical compound diagram. \n\nInternational Chemical Identifiers (\"InChI values\") are a standardized encoding to describe chemical compounds. They take the form of a string of letters, numbers and deliminators, often between 100 - 400 characters long. \n\nThe chemical diagrams are provided as PNG files, often of such low quality that it may take a human several seconds to decipher. \n\nLabel length and image quality become a serious challenge here, because we must predict labels for a very large quantity of images. There are 1.6 million images in the test set abd 2.4 million images available in the training set!","metadata":{}},{"cell_type":"code","source":"\"\"\"\n### note: load datasets before running this cell\n\n# Example (image, target label) pair\\n\\n'\nfor val in train_ds.unbatch().take(1):\n    print('Example Label:\\n', val['InChI'].numpy())\n    print('\\nCorresponding Image:', plt.imshow(val['image'][:,:,0], cmap='binary'))\n\"\"\"","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:40:12.312534Z","iopub.execute_input":"2021-07-04T18:40:12.313050Z","iopub.status.idle":"2021-07-04T18:40:12.328441Z","shell.execute_reply.started":"2021-07-04T18:40:12.312930Z","shell.execute_reply":"2021-07-04T18:40:12.327403Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## MODEL STRUCTURE: \n\n**Image CNN + Attention Features encoder --> text Attention + (optional )CNN feature layer decoder.**\n\nThis is a hybrid approach with:\n \n - Image Encoder from [*Show, Attend and Tell: Neural Image Caption Generation with Visual Attention*](https://proceedings.mlr.press/v37/xuc15.pdf).  Generate image feature vectors using intermediate layer outputs from a pretrained CNN. (Here I use the more modern EfficientNet model (recommended by [*Darien Schettler*](https://www.kaggle.com/dschettler8845/bms-efficientnetv2-tpu-e2e-pipeline-in-3hrs/notebook)) with fixed weights and a trainable Dense layer for customization.)\n \n - T2T encoder-decoder model from [*All You Need is Attention*](https://papers.nips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf) (Self-attention feature extraction for both encoder and decoder, joint encoder-decoder attention feature interactions, and a dense prediction output block. Includes parameters to independently set 0 or more encoder / decoder blocks.\n\n - ***PLUS*** *(optional):* Decoder Output Blocks placed in Series (not stacked). Increase the number of trainable parameters without adding inference computational complexity, while also allowing decoders to specialize on different regions of the output.\n \n - ***PLUS*** *(optional):* Is attention really all you need? Add a convolutional layer to enhance text features before decoder self-attention to experiment with performance differences with and without extra convolutional layer(s). Use of CNN's in NLP comes from [*Convolutional Sequence to Sequence Learning*](http://proceedings.mlr.press/v70/gehring17a.html.)\n\n - ***PLUS*** *(optional):* Beam-Search Alternative, an extra decoding layer applied after the full logits prediction has been made. This takes the form of a bidirectional RNN with attention, applied to the full logits sequence. Because a full (initial) prediction has already been made, computations can be parallelized using statefull RNNs. (See more details below.)\n\n----\n\n## NEXT STEPS:\n\n - Experiment with **\"Tokens-to-Token ViT\"** in place of the image CNN. (Technique from [*Training Vision Transformers from Scratch on ImageNet*](https://arxiv.org/pdf/2101.11986.pdf)\n  \n - Train my **Beam-search Alternative**. \n\n    - Beam search is a technique to modify model predictions to reflect the (local) maximum likelihood estimate. However, it is *very* local in that computation expense increases quickly with the number of character steps taken into account. This is also a hard-coded algorithm, which is somewhat contrary to the philosophy of deep learning.\n\n    - A *Beam-search Alternative* would be an extra decoding layer applied *after* the full logits prediction has been made. This might be in the form of a stateful, bidirectional RNN that is computationally parallizable because it is applied to the full logits sequence.\n\n    - Need to revamp code to accept main model changes made for TPU support.\n\n - Treat the number of convolutional layers (decoder feature extraction) and number of decoders places in series (decoder prediction output) as **new hyperparameters** to tune.\n\n----","metadata":{}},{"cell_type":"markdown","source":"## Contents\n\n1. [Imports](https://colab.research.google.com/drive/1i6LMwu7BRfs955U4AdtV2oaI_9_A_Awq#scrollTo=TjuUOVXao__C&line=4&uniqifier=1)\n2. [Data Pipeline](https://colab.research.google.com/drive/1i6LMwu7BRfs955U4AdtV2oaI_9_A_Awq#scrollTo=lrLHKs5Ni7Sz)\n3. [Model Layers](https://colab.research.google.com/drive/1i6LMwu7BRfs955U4AdtV2oaI_9_A_Awq#scrollTo=W0T-u0vZamI8)\n    - [InChI Encoding](https://colab.research.google.com/drive/1i6LMwu7BRfs955U4AdtV2oaI_9_A_Awq#scrollTo=DYApmA2lf1hp&line=1&uniqifier=1)\n    - [Image Encoding and Self-Attention](https://colab.research.google.com/drive/1i6LMwu7BRfs955U4AdtV2oaI_9_A_Awq#scrollTo=FESofcGdEaWF&line=1&uniqifier=1)\n    - [Decoder Self-Attention](https://colab.research.google.com/drive/1i6LMwu7BRfs955U4AdtV2oaI_9_A_Awq#scrollTo=6qFDs9RTjvod&line=1&uniqifier=1)\n    - [Joint Encoder-Decoder Attention](https://colab.research.google.com/drive/1i6LMwu7BRfs955U4AdtV2oaI_9_A_Awq#scrollTo=jP-t1MkKnD5L)\n    - [Decoder Head (Prediction Output)](https://colab.research.google.com/drive/1i6LMwu7BRfs955U4AdtV2oaI_9_A_Awq#scrollTo=38GA7wtNEhqW&line=1&uniqifier=1)\n    - [Update Mechanism](https://colab.research.google.com/drive/1i6LMwu7BRfs955U4AdtV2oaI_9_A_Awq#scrollTo=_2UR1DLljD0S&line=1&uniqifier=1)\n4. [Full Model](https://colab.research.google.com/drive/1i6LMwu7BRfs955U4AdtV2oaI_9_A_Awq#scrollTo=D6GIs3f3rpu0&line=1&uniqifier=1)\n5. [Training](https://colab.research.google.com/drive/1i6LMwu7BRfs955U4AdtV2oaI_9_A_Awq#scrollTo=otxdN02mf1ht&line=1&uniqifier=1)\n6. [Inference](https://colab.research.google.com/drive/1i6LMwu7BRfs955U4AdtV2oaI_9_A_Awq#scrollTo=Sbvzr5rdmjgs&line=5&uniqifier=1)\n\n---","metadata":{}},{"cell_type":"markdown","source":"## Contents\n\n1. Imports\n2. Data Pipeline\n3. Model Layers\n    - InChI Encoding\n    - Image Encoding and Self-Attention\n    - Decoder Self-Attention\n    - Joint Encoder-Decoder Attention\n    - Decoder Head (Prediction Output)\n    - Update Mechanism\n4. Full Model\n5. Training\n6. Inference\n\n---","metadata":{}},{"cell_type":"code","source":"#### PACKAGE IMPORTS ####\n\n# file management\nimport os\nos.environ['TF_ENABLE_ONEDNN.OPTS'] = '1'  # Intel's TF optimization\n\n# Tensorflow components\n!pip install -U -q tensorflow==2.5.0\nimport tensorflow as tf\nfrom tensorflow import keras\nimport tensorflow_hub as hub\nfrom tensorflow.data import TFRecordDataset\nfrom tensorflow.data.experimental import TFRecordWriter\n\n# Text processing\nimport re\nimport string\n\n# metric for Kaggle Competition\n!pip install -q leven\nfrom leven import levenshtein\n\n# Kaggle (for TPU)\nfrom kaggle_datasets import KaggleDatasets\n\n# Visualizations\nimport matplotlib.pyplot as plt\n%matplotlib inline\nfrom PIL import Image\n\n\"\"\" Use on Colab\n# extra imports for GCS\n!pip install -q fsspec\n!pip install -q gcsfs\nimport fsspec, gcsfs \n\n# Tensorboard Profiler\n!pip install -U -q tensorboard\n!pip install -U -q tensorboard_plugin_profile\n!pip install --upgrade -q \"cloud-tpu-profiler>=2.3.0\"\n%load_ext tensorboard\n\n\n# Debugger\ntf.debugging.experimental.enable_dump_debug_info('./logs/', tensor_debug_mode=\"FULL_HEALTH\", \n                                                 circular_buffer_size=-1)\n\"\"\"\n\n# data management\nimport numpy as np\nimport pandas as pd\nimport itertools","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:40:12.330079Z","iopub.execute_input":"2021-07-04T18:40:12.330454Z","iopub.status.idle":"2021-07-04T18:41:40.098626Z","shell.execute_reply.started":"2021-07-04T18:40:12.330419Z","shell.execute_reply":"2021-07-04T18:41:40.096769Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# TF Rec directory\nKaggleDatasets().get_gcs_path('bmsshards')","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:41:40.102407Z","iopub.execute_input":"2021-07-04T18:41:40.102823Z","iopub.status.idle":"2021-07-04T18:41:48.236429Z","shell.execute_reply.started":"2021-07-04T18:41:40.102782Z","shell.execute_reply":"2021-07-04T18:41:48.235221Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model parameters\n\nThe 'ModelParameters' class manages global hyperparamaters for portability between Colab and Kaggle notebook environments. Once set, all other cells will run on either platform.\n\nOn Colab, connection to my personal Google Drive is required, as ModelParameters will extract the dataset from a zip file to the hosted environment. This process may take several minutes. (It would not be difficult for the reader to update the code to point to their own drive and download the zip dataset using the Kaggle API code below.)","metadata":{}},{"cell_type":"code","source":"\"\"\" Kaggle api for download the compressed dataset from Kaggle's servers to Colab.\n\n# imports\n!pip uninstall -y kaggle\n!pip install --upgrade pip\n!pip install kaggle==1.5.6\n\n# if needed, download data using '!kaggle competitions download -c bms-molecular-translation'\n# then unzip with '! unzip bms-molecular-translation.zip -d datasets'\nos.environ['KAGGLE_CONFIG_DIR'] = '/content/gdrive/MyDrive/Kaggle'  # api token location\n\"\"\"","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:41:48.238562Z","iopub.execute_input":"2021-07-04T18:41:48.238916Z","iopub.status.idle":"2021-07-04T18:41:48.246477Z","shell.execute_reply.started":"2021-07-04T18:41:48.238882Z","shell.execute_reply":"2021-07-04T18:41:48.245388Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# check for TPU & initialize\ntry:\n    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')\n    tf.config.experimental_connect_to_cluster(resolver)\n\n    tf.tpu.experimental.initialize_tpu_system(resolver)\n    print(\"All devices: \", tf.config.list_logical_devices('TPU'))\n\n    TPU = True\n    STRATEGY = tf.distribute.TPUStrategy(resolver)\n    os.environ[\"TFHUB_MODEL_LOAD_FORMAT\"] = \"UNCOMPRESSED\"  # for TF Hub models on TPU\n\n    # set mixed precision type\n    PRECISION_TYPE = 'mixed_bfloat16' \n    #PRECISION_TYPE = 'float32' \n\nexcept:\n    TPU = False\n    STRATEGY = tf.distribute.get_strategy()\n    PRECISION_TYPE = 'mixed_float16'\n    \n# set mixed precision policy\ntf.keras.mixed_precision.set_global_policy(PRECISION_TYPE)","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:41:48.248110Z","iopub.execute_input":"2021-07-04T18:41:48.248433Z","iopub.status.idle":"2021-07-04T18:41:54.690236Z","shell.execute_reply.started":"2021-07-04T18:41:48.248403Z","shell.execute_reply":"2021-07-04T18:41:54.689079Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ModelParameters:\n    def __init__(self, cloud_server='kaggle'):\n               \n        # universal parameters\n        self._batch_size = 16  # used on GPU. TPU batch size increased below\n        self._padded_length = 200\n        self._image_size = (540, 540)  # shape to process images in data pipeline, matches HUB model\n        self.SOS_string = 'InChI=1S/'  # start of sentence value\n        self.EOS_string = '<EOS>'  # end of sentence value\n        self._strategy = STRATEGY\n        self._precision_type = PRECISION_TYPE\n        self._tpu = TPU\n        \n        # TPU batch size\n        if self._tpu:\n            self._batch_size = 128 * self._strategy.num_replicas_in_sync\n\n        # File Paths       \n        if cloud_server == 'colab':  # Google Colab\n            \n            # load drive for saving checkpoints\n            try:\n                from google.colab import drive\n                drive.mount('/content/gdrive/') \n            except:\n                pass  # drive already mounted\n            \n            # check for TPU \n            if self._tpu: \n\n                # TPU file structure (via Kaggle GCS folder)\n                self._dataset_dir = 'gs://kds-df3031ee4e277d641d1044cc3e9386a923ca98833b0d51a2575d2932' # from Kaggle. Get updated directory on Kaggle via KaggleDatasets().get_gcs_path('bms-molecular-translation')\n                self._prepared_files_dir = 'gs://kds-96b617b700ddb4d07bc42a47c0a7abfe3a68d4510c459b7cd7b216e6'  # from Kaggle. Get updated directory on Kaggle via KaggleDatasets().get_gcs_path('periodic-table')\n                self._tfrec_dir = 'gs://kds-dc74fe0494d010e8c9544cd7fff86e64f08cb0cffd4c608156ff3f41'  # from Kaggle. Get updated directory on Kaggle via KaggleDatasets().get_gcs_path('bmsshards')\n                self._checkpoint_dir = '/content/gdrive/MyDrive/Colab_Notebooks/models/MolecularTranslation/checkpoints/'  # gdrive\n                self._load_checkpoint_dir = self._checkpoint_dir\n                self._csv_save_dir = './'\n\n            else:\n                # unzip data\n                if not os.path.isdir('/content/bms-molecular-translation'):\n                    !unzip -q /content/gdrive/MyDrive/Colab_Notebooks/models/MolecularTranslation/bms-molecular-translation.zip -d '/content/bms-molecular-translation'\n                \n                # file paths\n                self._dataset_dir = 'bms-molecular-translation/'\n                self._prepared_files_dir = '/content/gdrive/MyDrive/Colab_Notebooks/models/MolecularTranslation/'\n                self._checkpoint_dir = '/content/gdrive/MyDrive/Colab_Notebooks/models/MolecularTranslation/checkpoints/'\n                self._load_checkpoint_dir = self._checkpoint_dir\n                self._csv_save_dir = self._prepared_files_dir \n                self._tfrec_dir = None\n                \n        elif cloud_server == 'kaggle': # Kaggle cloud notebook (CPU / GPU)\n            from kaggle_datasets import KaggleDatasets\n            \n            # check for TPU \n            if self._tpu: \n                \n                # file paths\n                self._dataset_dir = '' #note: GCSpath does working for this Kaggle directory\n                self._prepared_files_dir = KaggleDatasets().get_gcs_path('periodic-table')\n                self._tfrec_dir = KaggleDatasets().get_gcs_path('bmsshards')\n                self._checkpoint_dir = './'\n                self._load_checkpoint_dir = './'\n                self._csv_save_dir = './'\n\n            # set GPU instance info\n            else:  \n                # file paths\n                self._dataset_dir = '../input/bms-molecular-translation/'\n                self._prepared_files_dir = '../input/periodic-table/'\n                self._tfrec_dir = '../input/bmsshards/'\n                self._checkpoint_dir = './'\n                self._load_checkpoint_dir = '../input/k/mvenou/bms-molecular-translation/checkpoints/'\n                self._csv_save_dir = './'\n                self._tfrec_dir = None\n\n        # common file paths\n        self._periodic_table_csv = os.path.join(self._prepared_files_dir, 'periodic_table_elements.csv')\n        self._vocab_csv = os.path.join(self._prepared_files_dir, 'vocab.csv')        \n        self._test_images_dir = os.path.join(self._dataset_dir, 'test/')\n        self._train_images_dir = os.path.join(self._dataset_dir, 'train/')\n        self._extra_labels_csv = os.path.join(self._dataset_dir, 'extra_approved_InChIs.csv')\n        self._train_labels_csv = os.path.join(self._dataset_dir, 'train_labels.csv')\n        self._sample_submission_csv = os.path.join(self._dataset_dir, 'sample_submission.csv')\n        \n    # functions to access params\n    def padded_length(self):\n        return self._padded_length\n    def mixed_precision(self):\n        return self._precision_type\n    def tpu(self):\n        return self._tpu\n    def tfrec_dir(self):\n        return self._tfrec_dir\n    def cloud_server(self):\n        return self._cloud_server\n    def strategy(self):\n        return self._strategy\n    def csv_save_dir(self):\n        return self._csv_save_dir\n    def train_labels_csv(self):\n        return self._train_labels_csv\n    def vocab_csv(self):\n        return self._vocab_csv\n    def periodic_table_csv(self):\n        return self._periodic_table_csv\n    def batch_size(self):\n        return self._batch_size  \n    def image_size(self):\n        return self._image_size    \n    def SOS(self):\n        return self.SOS_string\n    def EOS(self):\n        return self.EOS_string\n    def train_images_dir(self):\n        return self._train_images_dir\n    def test_images_dir(self):\n        return self._test_images_dir   \n    def checkpoint_dir(self):\n        return self._checkpoint_dir\n    def load_checkpoint_dir(self):\n        return self._load_checkpoint_dir\n","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:41:54.692219Z","iopub.execute_input":"2021-07-04T18:41:54.692647Z","iopub.status.idle":"2021-07-04T18:41:54.731993Z","shell.execute_reply.started":"2021-07-04T18:41:54.692596Z","shell.execute_reply":"2021-07-04T18:41:54.730626Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Initialize Parameter Options","metadata":{}},{"cell_type":"code","source":"PARAMETERS = ModelParameters(cloud_server='kaggle')","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:41:54.734157Z","iopub.execute_input":"2021-07-04T18:41:54.734673Z","iopub.status.idle":"2021-07-04T18:42:34.418264Z","shell.execute_reply.started":"2021-07-04T18:41:54.734620Z","shell.execute_reply":"2021-07-04T18:42:34.417114Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Input Pipeline**","metadata":{}},{"cell_type":"code","source":"# Load CSV as dataframe\nif not PARAMETERS.tpu():\n    train_labels_df = pd.read_csv(PARAMETERS.train_labels_csv())\n    train_labels_df.head()","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.421590Z","iopub.execute_input":"2021-07-04T18:42:34.422053Z","iopub.status.idle":"2021-07-04T18:42:34.427820Z","shell.execute_reply.started":"2021-07-04T18:42:34.422004Z","shell.execute_reply":"2021-07-04T18:42:34.426719Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### InChI Text Parsing\n\nWe split each InChI label into its \"vocabulary\" of logical subunits, consisting of element abbreviations numbers, common symbols and the required string 'InChI=1S/', which is at the start of every InChI label. We want to narrow down this vocabulary to the smallest set represented in our training data. The functions below provide a system for finding this minimal set, as well as preparing a new CSV file with parsed labels ready to be fed into a tokenizer layer.\n\n(For clarity and to reduce reliance on loading external files, the true code has been commented out and replaced with corresponding hard-coded values.)","metadata":{}},{"cell_type":"code","source":"def inchi_parsing_regex(parameters=PARAMETERS):\n    # regex for spliting on InChi, but preserving chemical element abbreviations and three-digit numbers\n    \n    # shortcut: hard coded values\n    vocab = [parameters.EOS(), parameters.SOS(), '(',\n            ')', '+', ',', '-', '/', 'Br', 'B', 'Cl', 'C', 'D', 'F',\n            'H', 'I', 'N', 'O', 'P', 'Si', 'S', 'T', 'b', 'c', 'h', 'i',\n            'm', 's', 't']\n        \n    vocab += [str(num) for num in reversed(range(168))]\n    vocab = [re.escape(val) for val in vocab]\n       \n    \"\"\" # to create vocab from scratch, use:\n    SOS = parameters.SOS()\n    EOS = parameters.EOS()\n    \n    # load list of elements we should search for within InChI strings: \n    periodic_elements = pd.read_csv(PARAMETERS.periodic_table_csv(), header=None)[1].to_list()\n    periodic_elements = periodic_elements + [val.lower() for val in periodic_elements] + [val.upper() for val in periodic_elements]\n    \n    punctuation = list(string.punctuation)\n    punctuation = [re.escape(val) for val in punctuation]   # update values with regex escape chars added as needed\n\n    three_dig_nums_list = [str(i) for i in range(1000, -1, -1)]\n\n    vocab = [SOS, EOS] + periodic_elements + three_dig_nums_list + punctuation\n    \"\"\"\n\n    split_elements_regex = rf\"({'|'.join(vocab)})\"\n    \n    return split_elements_regex","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.430096Z","iopub.execute_input":"2021-07-04T18:42:34.430421Z","iopub.status.idle":"2021-07-04T18:42:34.442985Z","shell.execute_reply.started":"2021-07-04T18:42:34.430390Z","shell.execute_reply":"2021-07-04T18:42:34.441918Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"INCHI_PARSING_REGEX = inchi_parsing_regex()\n\ndef parse_InChI(texts, parsing_regex=INCHI_PARSING_REGEX):  \n    return ' '.join(re.findall(parsing_regex, texts))\n\n\n# TF dataset map-compatible version\ndef parse_InChI_py_fn(texts, parsing_regex=INCHI_PARSING_REGEX):\n    def tf_parse_InChI(texts):  \n        texts = np.char.array(texts.numpy())\n        texts = np.char.decode(texts).tolist()\n        texts = tf.constant([parse_InChI(val) for val in texts])\n        return tf.squeeze(texts)\n    return tf.py_function(func=tf_parse_InChI, inp=[texts], Tout=tf.string)\n\n\n# extracts filepath from image name\ndef path_from_image_id(x, root_folder):\n    folder_a = tf.strings.substr(x, pos=0, len=1)\n    folder_b = tf.strings.substr(x, pos=1, len=1)\n    folder_c = tf.strings.substr(x, pos=2, len=1)\n    filename =  tf.strings.join([x, '.png'])\n    return tf.strings.join([root_folder, folder_a, folder_b, folder_c, filename], separator='/')","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.444858Z","iopub.execute_input":"2021-07-04T18:42:34.445524Z","iopub.status.idle":"2021-07-04T18:42:34.461945Z","shell.execute_reply.started":"2021-07-04T18:42:34.445479Z","shell.execute_reply":"2021-07-04T18:42:34.461058Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Tokenizer\n\nNote: this is done during preprocessing b/c TPU doesn't work well with TextVectorizer","metadata":{}},{"cell_type":"code","source":"def Tokenizer(parameters):\n    \"\"\" note: crops /pads to max len\n    \"\"\"\n\n    SOS = parameters.SOS()\n    EOS = parameters.EOS()\n    padded_length = PARAMETERS.padded_length()\n    \n    # Create vocabulary for tokenizer\n    def create_vocab():       \n        hard_coded_vocab = [PARAMETERS.EOS(), PARAMETERS.SOS(), '(',\n            ')', '+', ',', '-', '/', 'B', 'Br',  'C', 'Cl', 'D', 'F',\n            'H', 'I', 'N', 'O', 'P', 'S', 'Si', 'T', 'b', 'c', 'h', 'i',\n            'm', 's', 't']\n        \n        numbers = [str(num) for num in range(168)]\n        \n        vocab = hard_coded_vocab + numbers\n        \n        \"\"\"\n        # get from saved file\n        vocab = pd.read_csv(PARAMETERS.vocab_csv())['vocab_value'].to_list()   \n        vocab = list(vocab)\n        \"\"\"\n\n        \"\"\" \n        # To create from scratch, extract all vocab elements appearing in train set:\n        df = pd.read_csv(PARAMETERS.train_labels_csv())  \n        seg_len = 250000\n        num_breaks = len(df) // seg_len\n\n        vocab = set()\n        for i in range(num_breaks):\n\n            df_i =  df['InChI'].iloc[seg_len * i: seg_len * (i+1)]\n            texts =  df_i.apply(lambda x: set(parse_InChI(x).split()))\n            texts = texts.tolist()\n\n            vocab = vocab.union(*texts)\n\n            print(f'completed {i} / {num_breaks}')\n\n        vocab = list(vocab)\n        vocab_df = pd.DataFrame({'vocab_value': vocab})\n\n        # save results\n        filename = os.path.join(PARAMETERS.csv_save_dir(), 'vocab.csv')\n        vocab_df.to_csv(filename, index=False)\n        \"\"\"\n               \n        return vocab\n\n    vocab = create_vocab()\n    \n    # create tokenizer\n    tokenizer_layer = tf.keras.layers.experimental.preprocessing.TextVectorization(\n        standardize=None, split=lambda x: tf.strings.split(x, sep=' ', maxsplit=-1), \n        output_mode='int', output_sequence_length=padded_length, vocabulary=vocab)\n\n    # record EOS token\n    tokenized_EOS = tokenizer_layer(tf.constant([EOS]))\n    \n    # create inverse (de-tokenizer)\n    inverse_tokenizer = tf.keras.layers.experimental.preprocessing.StringLookup(\n        vocabulary=tokenizer_layer.get_vocabulary(), invert=True)\n\n    return tokenizer_layer, inverse_tokenizer, tokenized_EOS","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.463714Z","iopub.execute_input":"2021-07-04T18:42:34.464299Z","iopub.status.idle":"2021-07-04T18:42:34.481506Z","shell.execute_reply.started":"2021-07-04T18:42:34.464261Z","shell.execute_reply":"2021-07-04T18:42:34.480307Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"TOKENIZER_LAYER, INVERSE_TOKENIZER, TOKENIZED_EOS = \\\n    Tokenizer(parameters=PARAMETERS)\n\ndef tokenize_text(w, x, y, z):\n    # note: requires batch dim\n    y = TOKENIZER_LAYER(y)\n    return w, x, y, z","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.482944Z","iopub.execute_input":"2021-07-04T18:42:34.483327Z","iopub.status.idle":"2021-07-04T18:42:34.550843Z","shell.execute_reply.started":"2021-07-04T18:42:34.483291Z","shell.execute_reply":"2021-07-04T18:42:34.549806Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Image Loader","metadata":{}},{"cell_type":"code","source":"# Image loaders\ndef load_image(image_path):\n    image_path = tf.squeeze(image_path)\n    image = keras.layers.Lambda(lambda x: tf.io.read_file(x))(image_path)\n    return image   \n\ndef decode_image(image, target_size):\n    image = keras.layers.Lambda(lambda x: tf.io.decode_image(x, channels=1, expand_animations=False))(image)\n    image = keras.layers.experimental.preprocessing.Resizing(*target_size)(image)\n    return image    ","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.552490Z","iopub.execute_input":"2021-07-04T18:42:34.552816Z","iopub.status.idle":"2021-07-04T18:42:34.559935Z","shell.execute_reply.started":"2021-07-04T18:42:34.552773Z","shell.execute_reply":"2021-07-04T18:42:34.558660Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Datasets\n\nHere we create efficient tf.data.Dataset train / validation / test sets.\n\nOut data pipeline will read our prepared CSV of (image filename, parsed InChI and standard InChI) tuples. (If this file is not found, it will be created from scratch. This may take several minutes)  Iterating through the list, it will load batches of corresponding images and labels.\n\nOur datasets contain the following information, accessible by dict keys: images, image_id, InChI, parsed_InChI. (The test set uses InChI = parsed_InChI = 'InChI=1S/', the known required stating value for any InChI code.)","metadata":{}},{"cell_type":"markdown","source":"## Dataset Creation","metadata":{}},{"cell_type":"code","source":"def data_generator(image_set, parameters=PARAMETERS, labels_df=None, decode_images=True):\n       \n    # get global params\n    batch_size = parameters.batch_size()\n    target_size = parameters.image_size()\n    SOS = parameters.SOS()\n    EOS = parameters.EOS()\n    \n    # dataset options\n    options = tf.data.Options()\n    options.experimental_optimization.autotune_buffers = True\n    options.experimental_optimization.map_vectorization.enabled = True\n    options.experimental_optimization.apply_default_optimizations = True\n        \n    # Train & Validation Datasets\n    if image_set in ['train', 'valid']:\n        root_folder = parameters.train_images_dir()  # train / valid images\n        valid_split = 0.10\n        \n        # load labels into memory as dataframe\n        if labels_df is None:\n            labels_df = pd.read_csv(parameters.train_labels_csv())\n\n        # test / train split\n        num_valid_samples = int(valid_split * len(labels_df))\n        train_df = labels_df.iloc[num_valid_samples: ]  # get train split\n        valid_df = labels_df.iloc[: num_valid_samples]  # get validation split\n\n        # shuffle\n        train_df = train_df.sample(frac=1)\n        valid_df = valid_df.sample(frac=1)\n\n        # load into datasets  # (image_id, InChI)\n        train_ds = tf.data.Dataset.from_tensor_slices(train_df.values)\n        valid_ds = tf.data.Dataset.from_tensor_slices(valid_df.values)\n\n        train_ds = train_ds.with_options(options)\n        valid_ds = valid_ds.with_options(options)\n\n        # update image paths  \n        def map_path(x):  # (image_path, image_id, InChI)\n            image_id = x[0]\n            image_path = path_from_image_id(image_id, root_folder)\n            return image_path, x[0], x[1]\n\n        train_ds = train_ds.map(map_path, num_parallel_calls=tf.data.AUTOTUNE)\n        valid_ds = valid_ds.map(map_path, num_parallel_calls=tf.data.AUTOTUNE)\n\n        def map_parse(x, y, z):  # (image_path, image_id, InChI)\n            parsed_InChI = parse_InChI_py_fn(z)\n            return x, y, parsed_InChI, z\n   \n        train_ds = train_ds.map(map_parse, num_parallel_calls=tf.data.AUTOTUNE)\n        valid_ds = valid_ds.map(map_parse, num_parallel_calls=tf.data.AUTOTUNE)\n                \n        # load images into dataset       \n        def open_images(w, x, y, z):\n            w = load_image(w)\n            return w, x, y, z\n        \n        train_ds = train_ds.map(open_images, num_parallel_calls=tf.data.AUTOTUNE)\n        valid_ds = valid_ds.map(open_images, num_parallel_calls=tf.data.AUTOTUNE)    \n\n        # PREFETCH dataset BEFORE decoding images\n        train_ds = train_ds.prefetch(tf.data.AUTOTUNE)\n        valid_ds = valid_ds.prefetch(tf.data.AUTOTUNE)\n\n        def decode(w, x, y, z):\n            w = decode_image(w, target_size)\n            return w, x, y, z\n\n        if decode_images:\n            train_ds = train_ds.map(decode, num_parallel_calls=tf.data.AUTOTUNE)\n            valid_ds = valid_ds.map(decode, num_parallel_calls=tf.data.AUTOTUNE)    \n\n        # BATCH dataset AFTER decoding images (required by tf.io)\n        # should batch before other pure TF Lambda layer ops\n        train_ds = train_ds.batch(batch_size, drop_remainder=True)\n        valid_ds = valid_ds.batch(batch_size, drop_remainder=True)\n        \n        # add extra \"EOS\" values to end of parsed inchi\n        def extend_EOS(w, x, y, z):\n            y = tf.strings.join([y, EOS, EOS, EOS, EOS, EOS], separator=' ')\n            y = tf.reshape(y, [-1])\n            return w, x, y, z\n\n        train_ds = train_ds.map(extend_EOS, num_parallel_calls=tf.data.AUTOTUNE)\n        valid_ds = valid_ds.map(extend_EOS, num_parallel_calls=tf.data.AUTOTUNE)\n\n        # Tokenize parsed_inchi.  Note: ds must be batched before this step (size=1 is ok) \n        train_ds = train_ds.map(tokenize_text, num_parallel_calls=tf.data.AUTOTUNE)\n        valid_ds = valid_ds.map(tokenize_text, num_parallel_calls=tf.data.AUTOTUNE)\n\n        # name the elements\n        def map_names(w, x, y, z):\n            return  {'image': w, 'image_id': x, 'tokenized_InChI': y, 'InChI': z}\n        \n        train_ds = train_ds.map(map_names, num_parallel_calls=tf.data.AUTOTUNE)\n        valid_ds = valid_ds.map(map_names, num_parallel_calls=tf.data.AUTOTUNE)\n        \n        return train_ds, valid_ds\n    \n    # Test Dataset\n    elif image_set == 'test':\n\n        # note: image resizing and batching done during this loading step\n        # other elements must be batched before combining\n        image_ds = tf.keras.preprocessing.image_dataset_from_directory(\n            directory=parameters.test_images_dir(), labels='inferred', label_mode=None,\n            class_names=None, color_mode='grayscale', batch_size=1, \n            image_size=target_size, shuffle=False, seed=None, validation_split=None, \n            subset=None, follow_links=False)\n\n        # set filenames as label and batch\n        image_id_ds = tf.data.Dataset.from_tensor_slices(image_ds.file_paths)\n        image_id_ds = image_id_ds.map(lambda x: tf.strings.split(x, os.path.sep)[-1],\n                                      num_parallel_calls=tf.data.AUTOTUNE)\n        \n        # prepare images for TF Records creations. \n        # Note: do this step AFTER filenames step\n        if decode_images is False:  \n            # convert image to raw byte string. Note: cannot have batch dim for encoding\n            image_ds = image_ds.unbatch()\n            image_ds = image_ds.map(lambda x: tf.cast(x, dtype=tf.uint16))\n            image_ds = image_ds.map(lambda image: tf.io.encode_png(image))\n            image_ds = image_ds.map(lambda image: tf.io.serialize_tensor(image))\n            \n        # dataset consisting solely of InChI start 'InChI=1S/'\n        inchi_ds = image_id_ds.map(lambda x: tf.constant(SOS, dtype=tf.string),\n                                   num_parallel_calls=tf.data.AUTOTUNE)\n        \n        # merge datasets\n        test_ds = tf.data.Dataset.zip((image_ds, image_id_ds, inchi_ds, inchi_ds))\n        \n        # prefetch\n        test_ds = test_ds.prefetch(tf.data.AUTOTUNE)\n        test_ds = test_ds.batch(batch_size)\n\n        # Tokenize parsed_inchi.  Note: ds must be batched before this step (size=1 is ok) \n        test_ds = test_ds.map(tokenize_text, num_parallel_calls=tf.data.AUTOTUNE)\n\n        # set key names\n        def map_names(w, x, y, z):\n            return  {'image': w, 'image_id': x, 'tokenized_InChI': y, 'InChI': z}\n        \n        test_ds = test_ds.map(map_names, num_parallel_calls=tf.data.AUTOTUNE)\n        \n        return test_ds","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.561725Z","iopub.execute_input":"2021-07-04T18:42:34.562052Z","iopub.status.idle":"2021-07-04T18:42:34.594314Z","shell.execute_reply.started":"2021-07-04T18:42:34.562019Z","shell.execute_reply":"2021-07-04T18:42:34.592996Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Create Test, Train and Validation Datasets","metadata":{}},{"cell_type":"code","source":"if not PARAMETERS.tpu():\n    train_ds, valid_ds = data_generator('train', parameters=PARAMETERS, labels_df=train_labels_df, decode_images=True)\n    #test_ds = data_generator('test', parameters=PARAMETERS, labels_df=None, decode_images=True)\n    \n    print('Train DS')\n    for val in train_ds.take(1):    \n        print('image:', val['image'].shape, 'image_id:', val['image_id'].shape, \n              'InChI:', val['InChI'].shape, 'tokenized_InChI:', val['tokenized_InChI'].shape)\n\n    # Example (image, target label) pair\\n\\n'\n    print('Example Label:\\n', val['InChI'].numpy()[0])\n    print('\\nCorresponding Image:', plt.imshow(val['image'][0, :,:,:], cmap='binary'))","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.595926Z","iopub.execute_input":"2021-07-04T18:42:34.596391Z","iopub.status.idle":"2021-07-04T18:42:34.612439Z","shell.execute_reply.started":"2021-07-04T18:42:34.596348Z","shell.execute_reply":"2021-07-04T18:42:34.611229Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Examine data shapes","metadata":{}},{"cell_type":"code","source":"if not PARAMETERS.tpu():\n\n    print('Train DS')\n    for val in train_ds.take(1):    \n        print('image:', val['image'].shape, 'image_id:', val['image_id'].shape, \n              'InChI:', val['InChI'].shape, 'tokenized_InChI:', val['tokenized_InChI'].shape)\n\n    print('\\nValidation DS')\n    for val in valid_ds.take(1):\n        print('image:', val['image'].shape, 'image_id:', val['image_id'].shape, \n              'InChI:', val['InChI'].shape, 'tokenized_InChI:', val['tokenized_InChI'].shape)\n\n    try:\n        print('\\nTest DS')\n        for val in test_ds.take(1):\n            print('image:', val['image'].shape, 'image_id:', val['image_id'].shape, \n                'InChI:', val['InChI'].shape, 'tokenized_InChI:', val['tokenized_InChI'].shape)\n    except:\n        pass","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.613907Z","iopub.execute_input":"2021-07-04T18:42:34.614289Z","iopub.status.idle":"2021-07-04T18:42:34.626671Z","shell.execute_reply.started":"2021-07-04T18:42:34.614253Z","shell.execute_reply":"2021-07-04T18:42:34.625511Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## TF Records (optional)\n\n(Only used for TPU support)","metadata":{}},{"cell_type":"code","source":"# Create TF Examples\ndef make_example(image, image_id, tokenized_InChI, InChI):\n    image_feature = tf.train.Feature(\n        bytes_list=tf.train.BytesList(value=[image.numpy()])  # image provided as raw bytestring\n    )\n    image_id_feature = tf.train.Feature(\n        bytes_list=tf.train.BytesList(value=[image_id.numpy()])\n    )\n    tokenized_InChI_feature = tf.train.Feature(\n        bytes_list=tf.train.BytesList(value=[tf.io.serialize_tensor(tokenized_InChI).numpy()])\n    )\n    InChI_feature = tf.train.Feature(\n        bytes_list=tf.train.BytesList(value=[InChI.numpy()])\n    )\n\n    features = tf.train.Features(feature={\n        'image': image_feature,\n        'image_id': image_id_feature,\n        'tokenized_InChI': tokenized_InChI_feature,\n        'InChI': InChI_feature\n    })\n    \n    example = tf.train.Example(features=features)\n\n    return example.SerializeToString()\n\n\ndef make_example_py_fn(image, image_id, InChI, tokenized_InChI):\n    return tf.py_function(func=make_example, \n                   inp=[image, image_id, InChI, tokenized_InChI], \n                   Tout=tf.string)\n\n\n# Decode TF Examples\ndef decode_example(example, parameters=PARAMETERS):        \n    feature_description = {'image': tf.io.FixedLenFeature([], tf.string),\n                           'image_id': tf.io.FixedLenFeature([], tf.string),\n                           'tokenized_InChI': tf.io.FixedLenFeature([], tf.string),\n                           'InChI': tf.io.FixedLenFeature([], tf.string)}\n    \n    values = tf.io.parse_single_example(example, feature_description)\n    \n    \n    values['image'] = decode_image(values['image'], parameters.image_size())\n    values['tokenized_InChI'] = tf.io.parse_tensor(values['tokenized_InChI'],\n                                                  out_type=tf.int64)\n    values['tokenized_InChI'] = tf.cast(values['tokenized_InChI'], tf.int32)\n    \n    return values","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.628216Z","iopub.execute_input":"2021-07-04T18:42:34.628637Z","iopub.status.idle":"2021-07-04T18:42:34.642712Z","shell.execute_reply.started":"2021-07-04T18:42:34.628604Z","shell.execute_reply":"2021-07-04T18:42:34.641621Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def serialized_dataset_gen(parameters=PARAMETERS, set_type='train', labels_df=None):\n    \n    if set_type == 'train':\n        train_ds, valid_ds = data_generator(image_set='train', \n                                            parameters=parameters, \n                                            labels_df=train_labels_df, \n                                            decode_images=False)  # output images as bytestrings\n\n        train_ds = train_ds.unbatch()\n        valid_ds = valid_ds.unbatch()\n\n        # Create TF Examples\n        train_ds = train_ds.map(lambda x: make_example_py_fn(x['image'], x['image_id'], x['tokenized_InChI'], x['InChI']), \n                                num_parallel_calls=tf.data.AUTOTUNE)\n        valid_ds = valid_ds.map(lambda x: make_example_py_fn(x['image'], x['image_id'], x['tokenized_InChI'], x['InChI']), \n                                num_parallel_calls=tf.data.AUTOTUNE)\n        \n        return train_ds, valid_ds\n    \n    else: #test_set:\n        test_ds = data_generator(image_set='test', \n                                 parameters=parameters, \n                                 labels_df=None, \n                                 decode_images=False)  # output images as bytestrings\n        \n        test_ds = test_ds.unbatch()\n            \n        # Create TF Examples\n        test_ds = test_ds.map(lambda x: make_example_py_fn(x['image'], x['image_id'], x['tokenized_InChI'], x['InChI']), \n                              num_parallel_calls=tf.data.AUTOTUNE)\n        \n        return test_ds","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.644371Z","iopub.execute_input":"2021-07-04T18:42:34.644791Z","iopub.status.idle":"2021-07-04T18:42:34.659801Z","shell.execute_reply.started":"2021-07-04T18:42:34.644742Z","shell.execute_reply":"2021-07-04T18:42:34.658882Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create TF Record Shards\n\"\"\"\nNOTE: Changes have been made to the other dataset pipeline functions. \nTest / Revise this for compatability before running.\n\"\"\"\ndef create_records(dataset, subset, num_shards):\n    \n    folder = subset + '_tfrec'\n    \n    if subset =='train':\n        num_samples = int(.9 * len(train_labels_df))    # test / valid split\n    elif subset == 'valid':\n        num_samples = int(.1 * len(train_labels_df))\n    else:\n        num_samples = 2000000\n\n    if not os.path.isdir(folder):\n        os.mkdir(folder)\n        \n    for shard_num in range(num_shards):\n        \n        filename = os.path.join(folder, f'{subset}_shard_{shard_num+1}')\n        try:\n            this_shard = dataset.skip(shard_num * num_samples//num_shards).take(num_samples//num_shards)\n        \n            print(f'Writing shard {shard_num+1}/{num_shards} to {filename}')\n            writer = tf.data.experimental.TFRecordWriter(filename)\n            writer.write(this_shard)\n        except:\n            break\n    return None \n    \n# Load dataset from saved TF Record Shards\ndef dataset_from_records(subset, parameters=PARAMETERS):\n\n    # optimizations\n    options = tf.data.Options()\n    options.experimental_optimization.autotune_buffers = True\n    options.experimental_optimization.map_vectorization.enabled = True\n    options.experimental_optimization.apply_default_optimizations = True\n\n    filepath = os.path.join(parameters.tfrec_dir(), \n                            subset + '_tfrec/*')\n\n    dataset = tf.data.Dataset.list_files(filepath)  # put all tf rec filenames in a ds\n    dataset = dataset.shuffle(10**6)\n \n    # merge the files\n    num_readers = parameters.strategy().num_replicas_in_sync\n    dataset = dataset.interleave(tf.data.TFRecordDataset,  \n                                 cycle_length=num_readers, block_length=1,\n                                 deterministic=False, num_parallel_calls=tf.data.AUTOTUNE)\n    \n    dataset = dataset.shuffle(10**6)\n    \n    # decode examples\n    dataset = dataset.map(decode_example, num_parallel_calls=tf.data.AUTOTUNE)\n\n    # note: tokenized InChI element spec needs help determining shape\n    for val in dataset.take(1):\n        padded_length = val['tokenized_InChI'].shape[-1]\n\n    # coerce unknown shape\n    dataset = dataset.map(lambda x: {'image':x['image'],\n                                     'image_id': x['image_id'],\n                                     'tokenized_InChI': tf.reshape(x['tokenized_InChI'], [padded_length]),\n                                     'InChI': x['InChI']},\n                          num_parallel_calls=tf.data.AUTOTUNE)  \n\n    dataset = dataset.batch(parameters.batch_size(), drop_remainder=True)\n    dataset = dataset.prefetch(tf.data.AUTOTUNE)\n        \n    return dataset","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.661188Z","iopub.execute_input":"2021-07-04T18:42:34.661712Z","iopub.status.idle":"2021-07-04T18:42:34.677751Z","shell.execute_reply.started":"2021-07-04T18:42:34.661677Z","shell.execute_reply":"2021-07-04T18:42:34.676915Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# To create TF_Records files\n# note: can take 8+ hours for train set alone!\n\"\"\"\nserial_train_ds, serial_valid_ds = serialized_dataset_gen(parameters=PARAMETERS, set_type='train', labels_df=train_labels_df)\nserial_test_ds = serialized_dataset_gen(parameters=PARAMETERS, set_type='test', labels_df=None)\n\"\"\"","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.679050Z","iopub.execute_input":"2021-07-04T18:42:34.679595Z","iopub.status.idle":"2021-07-04T18:42:34.696406Z","shell.execute_reply.started":"2021-07-04T18:42:34.679557Z","shell.execute_reply":"2021-07-04T18:42:34.695378Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Load TF Record shards into datasets","metadata":{}},{"cell_type":"code","source":"# IF USING TF_RECORDS:\n#NOTE: If no dataset loads, check if Kaggle GCS directories have changed. (This happends periodically)\nif PARAMETERS.tpu():\n    with PARAMETERS.strategy().scope(): \n        train_ds = dataset_from_records('train', parameters=PARAMETERS)\n        valid_ds = dataset_from_records('valid', parameters=PARAMETERS)\n\n    print('Train DS')\n    for val in train_ds.take(1):    \n        print('image:', val['image'].shape, 'image_id:', val['image_id'].shape, \n              'InChI:', val['InChI'].shape, 'tokenized_InChI:', val['tokenized_InChI'].shape)\n\n    print('\\nValidation DS')\n    for val in valid_ds.take(1):\n        print('image:', val['image'].shape, 'image_id:', val['image_id'].shape, \n              'InChI:', val['InChI'].shape, 'tokenized_InChI:', val['tokenized_InChI'].shape)\n\n    try:\n        print('\\nTest DS')\n        for val in test_ds.take(1):\n            print('image:', val['image'].shape, 'image_id:', val['image_id'].shape, \n                'InChI:', val['InChI'].shape, 'tokenized_InChI:', val['tokenized_InChI'].shape)\n    except:\n        pass","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:42:34.697855Z","iopub.execute_input":"2021-07-04T18:42:34.698337Z","iopub.status.idle":"2021-07-04T18:43:00.565375Z","shell.execute_reply.started":"2021-07-04T18:42:34.698300Z","shell.execute_reply":"2021-07-04T18:43:00.563600Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Model Layers**","metadata":{}},{"cell_type":"markdown","source":"## InChI Encoding\n\nTokenizer and Embedding to convert parsed InChI strings to tensors of numbers","metadata":{}},{"cell_type":"code","source":"def InChIEncoder(vocab_size, inchi_embedding_dim, num_chars):\n\n    EmbeddingLayer = tf.keras.layers.Embedding(input_dim=vocab_size, \n        output_dim=inchi_embedding_dim, mask_zero=False, input_length=num_chars)\n    \n    inchi = keras.layers.Input([num_chars], name='tokenized_inchi')\n    start_var = keras.layers.Input([1, inchi_embedding_dim], name='start_var')\n    pos_encoding = keras.layers.Input([num_chars, inchi_embedding_dim], name='positional_encoding')\n\n    inputs = [inchi, start_var, pos_encoding]\n\n    # embedding\n    inchi = inchi[:, :-1]  # drop last val\n    inchi = EmbeddingLayer(inchi)\n\n    # append start token\n    start_var = tf.cast(start_var, dtype=inchi.dtype)\n    inchi = keras.layers.Concatenate(axis=1)([start_var, inchi])  \n    \n    # add positional encoding\n    pos_encoding = tf.cast(pos_encoding, dtype=inchi.dtype)\n    inchi = keras.layers.Add()([inchi, pos_encoding])\n\n    outputs = [inchi]\n\n    return keras.Model(inputs, outputs, name='InChIEncoder')","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:00.568881Z","iopub.execute_input":"2021-07-04T18:43:00.569458Z","iopub.status.idle":"2021-07-04T18:43:00.584590Z","shell.execute_reply.started":"2021-07-04T18:43:00.569413Z","shell.execute_reply":"2021-07-04T18:43:00.582536Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"temp_inchi_enc = InChIEncoder(vocab_size=199, inchi_embedding_dim=512, num_chars=200)\ntemp_inchi_enc.summary()","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:00.591908Z","iopub.execute_input":"2021-07-04T18:43:00.592347Z","iopub.status.idle":"2021-07-04T18:43:00.694451Z","shell.execute_reply.started":"2021-07-04T18:43:00.592294Z","shell.execute_reply":"2021-07-04T18:43:00.692852Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Image Encoder\n\nFeature Extraction Step 1: Run the images through a pre-trained image network, extracting features as the output of an intermediate convolutional layer. [Technique from \"Show, Attend and Tell: Neural Image Caption Generation with Visual Attention,\" cited at the top of this notebook.]  A dense layer is added for transfer learning and to control the dimension of the attention mechanism used later.","metadata":{}},{"cell_type":"markdown","source":"Transfer Model","metadata":{}},{"cell_type":"code","source":"def ImageFeaturesExtractor(image_shape):\n    \n    # Note: temporarily disable mixed precision during load. (Model doesn't handle it properly)\n    tf.keras.mixed_precision.set_global_policy('float32')  # removed mixed precision\n\n    base_transfer_model = keras.applications.EfficientNetB2(\n                            include_top=False, \n                            weights=None,\n                            input_shape=(*image_shape[:2], 3))\n    tf.keras.mixed_precision.set_global_policy('float32')  # removed mixed precision\n\n    # revert to orig mixed precision policy\n    tf.keras.mixed_precision.set_global_policy(PARAMETERS.mixed_precision()) \n\n    model = keras.Model(inputs=base_transfer_model.inputs, \n                        outputs=base_transfer_model.get_layer('top_activation').output, \n                        name='ImageFeaturesExtractor')\n    \n    return model","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:00.696918Z","iopub.execute_input":"2021-07-04T18:43:00.697355Z","iopub.status.idle":"2021-07-04T18:43:00.705314Z","shell.execute_reply.started":"2021-07-04T18:43:00.697312Z","shell.execute_reply":"2021-07-04T18:43:00.703963Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#TransferModel((224,224,1)).summary()","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:00.707232Z","iopub.execute_input":"2021-07-04T18:43:00.707702Z","iopub.status.idle":"2021-07-04T18:43:00.720581Z","shell.execute_reply.started":"2021-07-04T18:43:00.707651Z","shell.execute_reply":"2021-07-04T18:43:00.719053Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def ImageEncoder(image_shape, encoder_dim, ImageFeaturesExtractor):\n\n    num_features = ImageFeaturesExtractor.output_shape[-3] * ImageFeaturesExtractor.output_shape[-2]\n\n    image = keras.layers.Input(image_shape, name='image')\n    pos_encoding = keras.layers.Input([num_features, encoder_dim], name='positional_encoding')\n\n    inputs = [image, pos_encoding]\n    \n    # get features\n    image = ImageFeaturesExtractor(image)\n\n    # update dims\n    image = keras.layers.Dense(encoder_dim)(image)\n    image = keras.layers.Dense(encoder_dim)(image)\n\n    # reshape as feature vectors\n    image = keras.layers.Reshape([-1, encoder_dim])(image)\n\n    # add positional encoding\n    pos_encoding = tf.cast(pos_encoding, dtype=image.dtype)\n    image = keras.layers.Add(name='add_positional_encoding')([image, pos_encoding])\n\n    outputs = [image]\n\n    return keras.Model(inputs, outputs, name='ImageEncoder')","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:00.722937Z","iopub.execute_input":"2021-07-04T18:43:00.723573Z","iopub.status.idle":"2021-07-04T18:43:00.739482Z","shell.execute_reply.started":"2021-07-04T18:43:00.723523Z","shell.execute_reply":"2021-07-04T18:43:00.737747Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"temp_image_shape=(320,320,3)\n\nImageEncoder(image_shape=temp_image_shape, \n             encoder_dim=256,\n             ImageFeaturesExtractor=ImageFeaturesExtractor(temp_image_shape)).summary()","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:00.741214Z","iopub.execute_input":"2021-07-04T18:43:00.741780Z","iopub.status.idle":"2021-07-04T18:43:06.294108Z","shell.execute_reply.started":"2021-07-04T18:43:00.741721Z","shell.execute_reply":"2021-07-04T18:43:06.292752Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Encoder Attention\n\nFeature Extraction Step 2: Now that we have basic feature vectors, we use self-attention to generate more complex features. This is the encoding step used in \"Attention is All You Need,\" cited above. ","metadata":{}},{"cell_type":"code","source":"def EncoderAttention(num_blocks, num_attention_heads, features_dim, num_encoder_features):\n\n    # Inputs\n    encoder_features = keras.layers.Input([num_encoder_features, features_dim], name='encoder_features')\n    inputs = [encoder_features]\n\n    for i in range(num_blocks):\n        # self attention block\n        AttentionLayer = keras.layers.MultiHeadAttention(num_heads=num_attention_heads, \n                                       key_dim = features_dim//num_attention_heads, \n                                       dropout=0.1, name=f'SelfAttention_{i}')\n            \n        attention_features = AttentionLayer(query=encoder_features, \n                                            value=encoder_features, \n                                            attention_mask=None)\n        \n        encoder_features = keras.layers.Add()([encoder_features, attention_features])\n        encoder_features = keras.layers.LayerNormalization()(encoder_features)\n\n        # feed forward block\n        dense_features = keras.layers.Dense(features_dim, activation='relu')(encoder_features)\n        dense_features = keras.layers.Dense(features_dim, activation=None)(dense_features)\n\n        dense_features = keras.layers.Dropout(rate=.1)(dense_features)\n        encoder_features = keras.layers.Add()([encoder_features, dense_features])\n        encoder_features = keras.layers.LayerNormalization()(encoder_features)\n\n    outputs = [encoder_features]\n\n    return keras.Model(inputs, outputs, name='EncoderAttention')","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:06.295860Z","iopub.execute_input":"2021-07-04T18:43:06.296250Z","iopub.status.idle":"2021-07-04T18:43:06.307665Z","shell.execute_reply.started":"2021-07-04T18:43:06.296209Z","shell.execute_reply":"2021-07-04T18:43:06.306165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"EncoderAttention(num_blocks=1, num_attention_heads=8, features_dim=256, num_encoder_features=100).summary()","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:06.309212Z","iopub.execute_input":"2021-07-04T18:43:06.309560Z","iopub.status.idle":"2021-07-04T18:43:06.542168Z","shell.execute_reply.started":"2021-07-04T18:43:06.309526Z","shell.execute_reply":"2021-07-04T18:43:06.540933Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Decoder & Encoder / Decoder Attention\n\nText Feature extraction + Encoder/Decoder Joint Attention interaction.\n\nWith use_covolutions set to False, this is the decoder self-attention feature-extraction step from \"Attention is All You Need,\" cited above (with learned positional encoding). \n\nIncludes an (optional) parameter to add a small convolutional layer for feature enhancement before the attention layer. This is included for experimentation / verification that attention really is all you need.\n","metadata":{}},{"cell_type":"code","source":"def DecoderAttention(num_blocks, num_attention_heads, num_steps, encoder_dim, decoder_dim, num_encoder_features):\n\n    encoder_features = keras.layers.Input([num_encoder_features, encoder_dim], name='encoder_features')\n    decoder_features = keras.layers.Input([num_steps, decoder_dim], name='decoder_features')\n    \n    inputs = [encoder_features, decoder_features]\n\n    # self-attention mask\n    ones = tf.ones([num_steps, num_steps])\n    self_attention_mask = tf.linalg.band_part(ones, -1, 0)\n\n    for i in range(num_blocks):\n\n        # regularization\n        decoder_features = keras.layers.SpatialDropout1D(rate=.05)(decoder_features)\n\n        # Self Attention Block\n        SelfAttentionLayer = keras.layers.MultiHeadAttention(\n                                    num_heads=num_attention_heads, \n                                    key_dim = decoder_dim//num_attention_heads, \n                                    dropout=0.1, \n                                    name=f'SelfAttention_{i}')\n        \n        attention_features = SelfAttentionLayer(query=decoder_features, \n                                                value=decoder_features, \n                                                attention_mask=self_attention_mask)\n        decoder_features = keras.layers.Add()([decoder_features, attention_features])\n        decoder_features = keras.layers.LayerNormalization()(decoder_features)\n\n        \n        # Joint Attention Block\n        JointAttentionLayer = keras.layers.MultiHeadAttention(\n                                    num_heads=num_attention_heads, \n                                    key_dim = decoder_dim//num_attention_heads, \n                                    dropout=0.1, \n                                    name=f'JointAttention_{i}')\n        \n        attention_features = JointAttentionLayer(query=decoder_features, \n                                                 value=encoder_features, \n                                                 attention_mask=None)\n        decoder_features = keras.layers.Add()([decoder_features, attention_features])\n        decoder_features = keras.layers.LayerNormalization()(decoder_features)\n\n\n        # Feed Forward Block                \n        dense_features = keras.layers.Dense(decoder_dim, activation='relu')(decoder_features)\n        dense_features = keras.layers.Dense(decoder_dim, activation=None)(dense_features)\n\n        dense_features = keras.layers.Dropout(rate=.1)(dense_features)\n        decoder_features = keras.layers.Add()([decoder_features, dense_features])\n        decoder_features = keras.layers.LayerNormalization()(decoder_features)\n\n    outputs = [decoder_features]\n\n    return keras.Model(inputs, outputs, name='DecoderAttention')","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:06.543947Z","iopub.execute_input":"2021-07-04T18:43:06.544419Z","iopub.status.idle":"2021-07-04T18:43:06.560339Z","shell.execute_reply.started":"2021-07-04T18:43:06.544371Z","shell.execute_reply":"2021-07-04T18:43:06.559087Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DecoderAttention(num_blocks=1, num_attention_heads=8, num_steps=200, \n                  encoder_dim=256, decoder_dim=512,num_encoder_features=100).summary()","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:06.562040Z","iopub.execute_input":"2021-07-04T18:43:06.562461Z","iopub.status.idle":"2021-07-04T18:43:06.893965Z","shell.execute_reply.started":"2021-07-04T18:43:06.562423Z","shell.execute_reply":"2021-07-04T18:43:06.892714Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Decoder Head (Prediction Output)\n\nThis is where we use what was learned in the encoder-decoder attention to output predicted labels. It is the prediction step from \"Attention is All You Need.\"","metadata":{}},{"cell_type":"code","source":"def DecoderHead(decoder_dim, vocab_size, num_decoder_features, dual_heads_split_step=None):\n   \n    # Inputs\n    decoder_features = keras.layers.Input([num_decoder_features, decoder_dim], name='decoder_features')\n    inputs = [decoder_features]\n\n    # Model\n    if dual_heads_split_step is None:\n        logits = keras.layers.Dense(vocab_size)(decoder_features)\n\n    else:\n        decoder_features_0 = decoder_features[:, :dual_heads_split_step, :]\n        decoder_features_1 = decoder_features[:, dual_heads_split_step:, :]\n\n        logits_0 = keras.layers.Dense(vocab_size)(decoder_features_0)\n        logits_1 = keras.layers.Dense(vocab_size)(decoder_features_1)\n\n        logits = keras.layers.Concatenate(axis=1)([logits_0, logits_1])\n    \n    probabilities = keras.layers.Softmax(dtype=tf.float32)(logits)\n\n    outputs = [probabilities]\n\n    return keras.Model(inputs, outputs, name='DecoderHead')","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:06.895695Z","iopub.execute_input":"2021-07-04T18:43:06.896148Z","iopub.status.idle":"2021-07-04T18:43:06.904880Z","shell.execute_reply.started":"2021-07-04T18:43:06.896085Z","shell.execute_reply":"2021-07-04T18:43:06.903864Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DecoderHead(decoder_dim=512, vocab_size=199, num_decoder_features=200, dual_heads_split_step=None).summary()","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:06.906466Z","iopub.execute_input":"2021-07-04T18:43:06.906929Z","iopub.status.idle":"2021-07-04T18:43:06.983309Z","shell.execute_reply.started":"2021-07-04T18:43:06.906872Z","shell.execute_reply":"2021-07-04T18:43:06.982051Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Update Mechanism (Optional)\n\n*Note: this is fully coded but I have not had time to train parameters with it. I leave that as a future opportunity for exploration.*\n\nNLP technicques typically output logits to find the highest likelhood token prediction. This can be improved to a (local) maximum likelihood selection using a \"beam step\" that ay override the initial prediction choice. \n\nThis layer is an alternative system for updating predictions. Unlike \"beam,\" it is trainable and includes longer-range dependencies (instead of the very \"local\" beam step.) The entire original prediction is passed through a bidirectional RNN. ","metadata":{}},{"cell_type":"code","source":"\"\"\"\ndef BeamUpdate(num_beam_blocks, num_att_blocks,\n               num_encoder_vectors, encoder_units, \n               decoder_units, max_len, vocab_size, name='BeamUpdate'):\n    \n    # update to required GRU model dtypes\n    tf.keras.mixed_precision.set_global_policy('float32')\n    \n    # layers\n    # note: GRU doesn't appear to be compatible with reduced precision\n    BeamUnit = keras.layers.GRU(decoder_units, return_sequences=True, \n                    return_state=True, go_backwards=True,\n                    dtype=tf.keras.mixed_precision.Policy('float32'))  \n\n    use_convolutions = False\n    BeamDecoderAttention = DecoderAttention(num_att_blocks, encoder_units, decoder_units, \n                                            num_encoder_vectors, max_len, \n                                            use_convolutions, name='BeamDecoderAttention')\n\n    BeamDecoderHead = DecoderHead(decoder_units, vocab_size,  \n                                  use_dual_heads=False, split_char_num=None,\n                                  name='BeamDecoderHead')\n\n    \n    # Inputs\n    beam_input = keras.layers.Input([max_len, vocab_size], name='beam_input') \n    hidden_state = keras.layers.Input([decoder_units], name='hidden_state')\n    encoder_features = keras.layers.Input([num_encoder_vectors, encoder_units], name='encoder_features')   # from image \n    mask = keras.layers.Input([max_len, max_len], name='mask')   # should pass in all 1's, i.e. no masking\n    \n    inputs = [beam_input, hidden_state, encoder_features, mask]\n\n    # create initial hidden state for RNN dim = decoder_units\n    beam_hidden_state = tf.reduce_mean(encoder_features, -2)\n    beam_hidden_state = keras.layers.Dense(decoder_units, activation='relu')(beam_hidden_state)\n    \n    # Decoder encoding using 1 or more Beam layers\n    for i in range(num_beam_blocks):\n        beam_out, beam_hidden_state = \\\n            BeamUnit(beam_input, initial_state=[beam_hidden_state])\n\n    # Attention & Prediction (uses \"Attention is All You Need\" structure)\n    decoder_features = BeamDecoderAttention([encoder_features, beam_out, mask])\n\n    probs = BeamDecoderHead([decoder_features])  \n\n    outputs = [probs]\n\n    return keras.Model(inputs, outputs, name=name)\n\n\nBeamUpdate(num_beam_blocks=1, num_att_blocks=1, num_encoder_vectors=50, encoder_units=256, \n            decoder_units=512, max_len=200, vocab_size=160, name='BeamUpdate').summary()\n\"\"\"","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:06.984828Z","iopub.execute_input":"2021-07-04T18:43:06.985204Z","iopub.status.idle":"2021-07-04T18:43:06.994923Z","shell.execute_reply.started":"2021-07-04T18:43:06.985158Z","shell.execute_reply":"2021-07-04T18:43:06.993737Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Full Model**\n\nAll the components are combined into a full encoder/decoder model. This is implemented using the subclassing API with custom call, train,  evaluation and prediction steps. Once initialized, the models have full access to high-level model.fit(), model.compile() and model.save_weights() methods.\n\nAn extra features implemented is having Decoder() elements in *series* (not stacked). This adds more trainable parameters without affecting inference speed, and allows decoders to specialize more on different regions of the text.\n\nBaseTrainer() model has the BeamUpdate mechanism disabled. InchiGenerator() models include the BeamUpdate.","metadata":{}},{"cell_type":"code","source":"class TrainerModel(keras.Model):\n\n    def __init__(self, encoder_blocks, encoder_dim, decoder_blocks, decoder_dim, \n                 dual_heads_split_step=None, parameters=PARAMETERS, \n                 name='TrainerModel', **kwargs):\n        \n        super().__init__(name=name, **kwargs)\n\n        self.encoder_blocks = encoder_blocks\n        self.encoder_dim = encoder_dim\n        self.decoder_blocks = decoder_blocks\n        self.decoder_dim = decoder_dim\n        self.dual_heads_split_step = dual_heads_split_step\n        self.parameters = parameters\n\n        tokenizer_layer, self.inverse_tokenizer, self.tokenized_EOS = \\\n            Tokenizer(parameters=self.parameters)\n        self.vocab_size = tokenizer_layer.vocabulary_size()\n        self.EOS = parameters.EOS()\n\n\n    def get_config(self):\n        config = {'encoder_blocks': self.encoder_blocks,\n                  'encoder_dim': self.encoder_dim,\n                  'decoder_blocks':self.decoder_blocks,\n                  'decoder_dim': self.decoder_dim,\n                  'dual_heads_split_step': dual_heads_split_step,\n                  'parameters':self.parameters,\n        }\n        return config \n\n\n    def build(self, input_shape):\n\n        self.batch_size = input_shape[0][0]\n        image_shape = input_shape[0][1:]\n        tokenized_inchi_shape = input_shape[1][1:]\n        self.padded_length = tokenized_inchi_shape[0]\n        \n        ###### InChI  ######\n        # InChI start variable\n        initializer = tf.random_normal_initializer()(shape=[1, 1, self.decoder_dim])\n        start_var = tf.Variable(initializer, trainable=True, dtype=tf.float32, name='start_var')\n        self.start_var = tf.tile(start_var, [self.batch_size, 1, 1])\n\n        # InChI encoder\n        self.InChIEncoder = InChIEncoder(vocab_size=self.vocab_size, \n                                         inchi_embedding_dim=self.decoder_dim, \n                                         num_chars=self.padded_length)\n        \n        # InChI positional encoding variable\n        initializer = tf.random_normal_initializer()(\n            shape=[1, self.padded_length, self.decoder_dim], dtype=tf.float32)\n        positional_encoding_inchi = tf.Variable(initializer, trainable=True, \n                                                     name='positional_encoding_inchi')\n        self.positional_encoding_inchi = tf.tile(positional_encoding_inchi, \n                                                 [self.batch_size, 1, 1])\n\n        ###### Image  ######\n        self.ImageFeaturesExtractor = ImageFeaturesExtractor(image_shape)\n        \n        self.ImageEncoder = ImageEncoder(image_shape=image_shape, \n                                         encoder_dim=self.encoder_dim, \n                                         ImageFeaturesExtractor=self.ImageFeaturesExtractor)\n        \n        self.num_image_features = self.ImageEncoder.output_shape[-2]\n\n        # Image positional encoding variable\n        initializer = tf.random_normal_initializer()(\n            shape=[1, self.num_image_features, self.encoder_dim], dtype=tf.float32)\n        positional_encoding_image = tf.Variable(initializer, trainable=True, \n                                                     name='positional_encoding_image')\n        self.positional_encoding_image = tf.tile(positional_encoding_image, \n                                                 [self.batch_size, 1, 1])\n\n        ###### Transformers  ######\n        self.EncoderAttention = EncoderAttention(num_blocks=self.encoder_blocks, \n                                                 num_attention_heads=8, \n                                                 features_dim=self.encoder_dim,\n                                                 num_encoder_features=self.num_image_features)\n        \n        self.DecoderAttention = DecoderAttention(num_blocks=self.decoder_blocks, \n                                                 num_steps=self.padded_length,\n                                                 num_attention_heads=8,\n                                                 encoder_dim=self.encoder_dim,\n                                                 decoder_dim=self.decoder_dim,\n                                                 num_encoder_features=self.num_image_features)\n        \n        num_decoder_features = self.DecoderAttention.output_shape[1]\n        \n        self.DecoderHead = DecoderHead(decoder_dim=self.decoder_dim, \n                                       vocab_size=self.vocab_size, \n                                       num_decoder_features=num_decoder_features, \n                                       dual_heads_split_step=self.dual_heads_split_step)\n        \n    def call(self, inputs, training=False):\n        \n        image = inputs[0]\n        tokenized_inchi = inputs[1]\n\n        # encoder\n        encoder_features = self.ImageEncoder(\n            [image, self.positional_encoding_image], training=training)\n\n        encoder_features = self.EncoderAttention(encoder_features, training=training)\n\n        # decoder\n        decoder_features = self.InChIEncoder(\n            [tokenized_inchi, self.start_var, self.positional_encoding_inchi], training=training)\n\n        decoder_features = self.DecoderAttention(\n            [encoder_features, decoder_features], training=training)\n\n        # predictions\n        probabilities = self.DecoderHead(decoder_features, training=training)\n\n        return probabilities\n\n    # define XLA compatible generation loop\n    def xla_generation_loop(self, inputs, training=False):\n        \n        image = inputs[0]\n        tokenized_inchi = inputs[1]\n\n        # get shapes\n        padded_length = self.padded_length\n        batch_size = self.batch_size\n\n        # encoder\n        encoder_features = self.ImageEncoder(\n            [image, self.positional_encoding_image], training=training)\n        \n        encoder_features = self.EncoderAttention(encoder_features, training=training)\n        \n        # decoder\n        # create containers\n        generated_probs = tf.TensorArray(dtype=tf.float32, size=padded_length)\n        generated_inchi = tf.TensorArray(dtype=tf.int32, size=padded_length)\n        \n        # initialize generated InChI values array\n        zeros = tf.zeros((self.batch_size, self.vocab_size), dtype=generated_probs.dtype)\n        for i in range(padded_length):\n            generated_probs = generated_probs.write(i, zeros)\n\n        # initialize generated InChI values array\n        zeros = tf.zeros((self.batch_size), dtype=generated_inchi.dtype)\n        for i in range(padded_length):\n            generated_inchi = generated_inchi.write(i, zeros)\n\n        # initialize step\n        step = tf.constant(0, dtype=tf.int32)\n        \n        # loop body function\n        def body_fn(generated_inchi, generated_probs, step):\n\n            inchi = tf.transpose(generated_inchi.stack(), [1,0])\n\n            decoder_features = self.InChIEncoder(\n                [inchi, self.start_var, self.positional_encoding_inchi], training=training)\n\n            decoder_features = self.DecoderAttention(\n                [encoder_features, decoder_features], training=training)\n\n            # get probabilities\n            probs = self.DecoderHead(decoder_features, training=training)     \n            \n            # select current step and save result\n            probs = probs[:, step, :]  \n            generated_probs = generated_probs.write(\n                step, tf.cast(probs, dtype=generated_probs.dtype))\n\n            # get new token prediction and save result\n            predicted_token = tf.argmax(probs, axis=-1)\n            generated_inchi = generated_inchi.write(\n                    step, tf.cast(predicted_token, dtype=generated_inchi.dtype))\n\n            # update step\n            step = step + 1\n            step = tf.cast(step, dtype=step.dtype)\n            \n            return [generated_inchi, generated_probs, step]\n\n        # loop conditional function\n        def cond_fn(generated_inchi, generated_probs, step):\n            return tf.math.less(step, padded_length)\n\n        # run generation loop\n        generated_inchi, generated_probs, step = \\\n            tf.while_loop(cond=cond_fn,\n                          body=body_fn,\n                          loop_vars=[generated_inchi, generated_probs, step],\n                          parallel_iterations=1,\n                          maximum_iterations=padded_length,\n                          shape_invariants=[None, None, tf.TensorShape([])],\n                          )\n        \n        # unpack generated probabilities\n        probabilities = tf.transpose(generated_probs.stack(), [1, 0, 2])\n\n        return probabilities\n    \n    # define quicker (but not TPU training compatible) generation loop\n    @tf.function(jit_compile=False, experimental_relax_shapes=True)\n    def generation_loop(self, inputs):\n\n        image = inputs[0]\n        tokenized_inchi = inputs[1]\n\n        # get shapes\n        padded_length = self.padded_length\n        batch_size = self.batch_size\n\n        # encoder\n        encoder_features = self.ImageEncoder(\n            [image, self.positional_encoding_image])\n        \n        encoder_features = self.EncoderAttention(encoder_features)\n        \n        # initialize token container\n        generated_inchi = tf.experimental.numpy.empty(\n            shape=[self.batch_size, 0], dtype=tokenized_inchi.dtype)\n\n        # token generation loop\n        for step in range(padded_length):\n\n            # zero pad future steps\n            inchi = tf.pad(generated_inchi, [[0,0],[0, padded_length - step]])\n\n            # decoder\n            decoder_features = self.InChIEncoder(\n                [inchi, self.start_var, self.positional_encoding_inchi])\n\n            decoder_features = self.DecoderAttention(\n                [encoder_features, decoder_features])\n\n            # get probabilities\n            probabilities = self.DecoderHead(decoder_features)\n\n            # get token predictions\n            predicted_token = tf.argmax(probabilities[:, step:step+1, :], axis=-1)\n            \n            # save results\n            generated_inchi = tf.concat(\n                [generated_inchi, tf.cast(predicted_token, dtype=generated_inchi.dtype)], axis=1)\n\n        return probabilities\n\n    @tf.function(jit_compile=False, experimental_relax_shapes=True)\n    def tokens_to_string(self, token_predictions):\n\n        # convert to strings\n        parsed_string_vals = self.inverse_tokenizer(token_predictions)\n        string_vals = keras.layers.Lambda(\n            lambda x: tf.strings.reduce_join(x, axis=-1))(parsed_string_vals)\n\n        # remove first EOS generated and everything after\n        pattern = ''.join([self.EOS, '.*$'])\n        string_vals = tf.strings.regex_replace(string_vals, pattern, rewrite='', \n                                               replace_global=True, name='remove_EOS')   \n\n        return string_vals","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:06.996585Z","iopub.execute_input":"2021-07-04T18:43:06.996940Z","iopub.status.idle":"2021-07-04T18:43:07.043878Z","shell.execute_reply.started":"2021-07-04T18:43:06.996902Z","shell.execute_reply":"2021-07-04T18:43:07.042785Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class EditDistanceMetric(tf.keras.metrics.Metric):\n    def __init__(self, name='edit_distance', **kwargs):\n        super().__init__(name=name, **kwargs)\n        self.edit_distance = self.add_weight(name='edit_distance', initializer='zeros')\n        self.batch_counter = self.add_weight(name='batch_counter', initializer='zeros')\n    \n    def update_state(self, y_true, y_pred, sample_weight=None):\n        y_true = tf.sparse.from_dense(y_true)\n        y_pred = tf.sparse.from_dense(tf.argmax(y_pred, axis=-1))  # convert probs to preds\n\n        y_true = tf.cast(y_true, tf.int32)\n        y_pred = tf.cast(y_pred, tf.int32)\n\n        # compute edit distance (of parsed tokens)\n        edit_distance = tf.edit_distance(y_pred, y_true, normalize=False)\n        self.edit_distance.assign_add(tf.reduce_mean(edit_distance))\n\n        # update counter\n        self.batch_counter.assign_add(tf.reduce_sum(1.))\n    \n    def result(self):\n        return self.edit_distance / self.batch_counter\n\n    def reset_state(self):\n        # The state of the metric will be reset at the start of each epoch.\n        self.edit_distance.assign(0.0)\n        self.batch_counter.assign(0.0)","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:07.045807Z","iopub.execute_input":"2021-07-04T18:43:07.046321Z","iopub.status.idle":"2021-07-04T18:43:07.067370Z","shell.execute_reply.started":"2021-07-04T18:43:07.046267Z","shell.execute_reply":"2021-07-04T18:43:07.066278Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Modified \"Attention is All You Need\" learning scheduler (to become cyclic)\nclass LRScheduleAIAYN(tf.keras.optimizers.schedules.LearningRateSchedule):\n\n    def __init__(self, scale_factor=1, warmup_steps=4000):  # defaults reflect paper's values\n        # cast dtypes\n        self.warmup_steps = tf.constant(warmup_steps, dtype=tf.float32)\n        dim = tf.constant(352, dtype=tf.float32)\n        scale_factor = tf.constant(scale_factor, dtype=tf.float32)\n        \n        self.scale = scale_factor * tf.math.pow(dim, -1.5)\n\n    def __call__(self, step):\n        step = tf.cast(step, tf.float32)\n        crit = self.warmup_steps\n\n        def false_fn(step):\n            adj_step = (step - crit) % (2.0*crit) + crit\n            return tf.math.pow(adj_step, -.5)\n\n        val = tf.cond(tf.math.less(step, crit),\n                      lambda: step * tf.math.pow(crit, -1.5),  # linear increase\n                      lambda: false_fn(step)  # decay\n                      )\n\n        return self.scale * val\n        \n    \"\"\"  # call for official AIAYN schedule\n    def __call__(self, step):\n        step = tf.cast(step, tf.float32)\n        if (step % self.warmup_steps) % 2 == 0:\n            val = step * tf.math.pow(self.warmup_steps, -1.5)  # linear increase\n        else:\n            val = tf.math.pow(step, -.5) # decay\n        return self.scale * val\n    \"\"\"\n\n\"\"\"\n# visualize learning rate \ntemp_lr = LRScheduleAIAYN()\nplt.plot([i for i in range(1, 16000)], [temp_lr(i) for i in range(1, 16000)])\nprint('Learning Rate Schedule')\n\"\"\"","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:07.069342Z","iopub.execute_input":"2021-07-04T18:43:07.069697Z","iopub.status.idle":"2021-07-04T18:43:07.094332Z","shell.execute_reply.started":"2021-07-04T18:43:07.069663Z","shell.execute_reply":"2021-07-04T18:43:07.093217Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build Model","metadata":{}},{"cell_type":"markdown","source":"Model compile options","metadata":{}},{"cell_type":"code","source":"NAME_MODIFIER = ''\n\n# build model\nIMAGE_DENSE_OUTPUT_DIM = 256  # note: only used with USE_DENSE_ENCODER_TOP = True.\nENCODER_ATT_UNITS = 256\nDECODER_UNITS = 512  # # \"All You Need is Attention\" uses 512 units\nBEAM_RNN_UNITS = 128  # note: only used in beam_model.\n\n# Note: model has capacity for up to 6 encoder and 6 decoder blocks. (as in AISAYN base model)\nNUM_ENCODER_BLOCKS = 0  # note: can set to 0 to skip encoder block\nNUM_DECODER_BLOCKS = 2  # max 6 enc and 6 dec with Colab memory constraints and cuts elsewhere\nUSE_DUAL_DECODERS = False\nUSE_CONVOLUTIONS = False\nif USE_CONVOLUTIONS:\n    checkpoint_save_name = 'ConvAtt_model_checkpoints' + NAME_MODIFIER\nelse:\n    checkpoint_save_name = 'AISAYN_model_checkpoints' + NAME_MODIFIER\n\nLOAD_CHECKPOINT_FILE = os.path.join(PARAMETERS.load_checkpoint_dir(), checkpoint_save_name, checkpoint_save_name)\nSAVE_CHECKPOINT_FILE = os.path.join(PARAMETERS.checkpoint_dir(), checkpoint_save_name, checkpoint_save_name)\n\n# note: in Kaggle,\n# LOAD_CHECKPOINT_FILE points to saved outputs from prev session\n# SAVE_CHECKPOINT_FILE points to saved outputs from current session","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:07.096256Z","iopub.execute_input":"2021-07-04T18:43:07.096835Z","iopub.status.idle":"2021-07-04T18:43:07.113198Z","shell.execute_reply.started":"2021-07-04T18:43:07.096782Z","shell.execute_reply":"2021-07-04T18:43:07.112148Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Initialize model","metadata":{}},{"cell_type":"code","source":"# Update inputs: remove string keys, as they are not compatible with TPU\ntrain_ds_int_index = train_ds.map(lambda x: (x['image'], x['tokenized_InChI'], \n                                            x['image_id'], x['InChI'])).prefetch(tf.data.AUTOTUNE)\nvalid_ds_int_index = valid_ds.map(lambda x: (x['image'], x['tokenized_InChI'], \n                                            x['image_id'], x['InChI'])).prefetch(tf.data.AUTOTUNE)","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:07.115108Z","iopub.execute_input":"2021-07-04T18:43:07.115631Z","iopub.status.idle":"2021-07-04T18:43:07.201891Z","shell.execute_reply.started":"2021-07-04T18:43:07.115578Z","shell.execute_reply":"2021-07-04T18:43:07.200757Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# callbacks\ncheckpoint = tf.keras.callbacks.ModelCheckpoint(SAVE_CHECKPOINT_FILE, monitor='loss', \n        save_weights_only=True, save_best_only=False, save_freq='epoch',\n        options=tf.train.CheckpointOptions(experimental_io_device='/job:localhost'))\n\nnan_stop = tf.keras.callbacks.TerminateOnNaN()","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:07.203638Z","iopub.execute_input":"2021-07-04T18:43:07.204076Z","iopub.status.idle":"2021-07-04T18:43:07.210780Z","shell.execute_reply.started":"2021-07-04T18:43:07.204036Z","shell.execute_reply":"2021-07-04T18:43:07.209018Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def compile_model(model=None, load_checkpoint=True, lr_scale_factor=1.0, label_smoothing=.1):\n\n    # compile using distribution strategy\n    with PARAMETERS.strategy().scope():\n     \n        # initialize if no model provided\n        if model is None:\n\n            model = TrainerModel(encoder_blocks=NUM_ENCODER_BLOCKS, \n                               encoder_dim=ENCODER_ATT_UNITS, \n                               decoder_blocks=NUM_DECODER_BLOCKS, \n                               decoder_dim=DECODER_UNITS, \n                               dual_heads_split_step=USE_DUAL_DECODERS,\n                               parameters=PARAMETERS)\n\n            if PARAMETERS.tpu():  \n        \n                temp_ds = PARAMETERS.strategy().experimental_distribute_dataset(train_ds_int_index)\n                temp_ds = iter(temp_ds)\n                val = next(temp_ds)\n\n                # build with new val (inference and training modes)\n                temp_func_train = tf.function(func=lambda x: model(x, True), experimental_relax_shapes=True,\n                                        experimental_follow_type_hints=True)\n                temp_func_inference = tf.function(func=lambda x: model(x, False), experimental_relax_shapes=True,\n                                        experimental_follow_type_hints=True)\n\n                PARAMETERS.strategy().run(temp_func_train, args=[(val[0], val[1])])  # use strategy.run() on TPU\n                PARAMETERS.strategy().run(temp_func_inference, args=[(val[0], val[1])])  # use strategy.run() on TPU\n                    \n            else:  \n                # build with original val\n                for val in train_ds_int_index.take(1): \n                    model(val, training=False)\n                    model(val, training=True)\n\n            # show summary\n            print(model.summary())\n            print('Models initialized.')\n        \n        # compiler components\n        # cyclic modification to AIAYN lr\n        learning_rate = .01 #LRScheduleAIAYN(scale_factor=lr_scale_factor, warmup_steps=5000)  \n\n        optimizer = tf.keras.optimizers.Adam(learning_rate,  # params from AIAYN\n                                             beta_1=0.9, beta_2=0.98, epsilon=10e-9)\n        \n        # metrics\n        edit_dist_metric = EditDistanceMetric()\n        \n        # loss with label smoothing\n        loss_fn = keras.losses.CategoricalCrossentropy(label_smoothing=label_smoothing)\n\n        # optimizations\n        tf.config.optimizer.set_jit(\"autoclustering\")  # XLA compiler optimization\n\n        if not PARAMETERS.tpu():\n            os.environ['TF_GPU_THREAD_MODE'] = 'gpu_private'  # better balances CPU / GPU interaction in tf.data    \n            optimizer = tf.keras.mixed_precision.LossScaleOptimizer(optimizer)  # required with mixed precision on GPU / CPU. Not used on TPU\n\n            # compile       \n            model.compile(optimizer=optimizer, \n                            loss=loss_fn,\n                            metrics=['categorical_accuracy', edit_dist_metric],\n                            steps_per_execution=8)\n        else:\n            # compile (note: EditDistance metric not compatible with TPU)\n            model.compile(optimizer=optimizer, \n                          loss=loss_fn,\n                          metrics=['categorical_accuracy'],\n                          steps_per_execution=8*PARAMETERS.strategy().num_replicas_in_sync)\n\n    if load_checkpoint:\n        # verify model calls & methods work\n        if not PARAMETERS.tpu():\n            model(val, training=False)\n\n            # sync weights\n            # WARNING!: in Kaggle this loads from prev session saved weights\n            try:\n                with PARAMETERS.strategy().scope(): \n                    model.load_weights(LOAD_CHECKPOINT_FILE)  \n                    pass\n            except:\n                print('No weights loaded')  \n\n        else:\n            # sync weights\n            # WARNING!: in Kaggle this loads from prev session saved weights\n            try:\n                with PARAMETERS.strategy().scope(): \n                    model.load_weights(LOAD_CHECKPOINT_FILE, \n                                            options=tf.train.CheckpointOptions(experimental_io_device=\"/job:localhost\"))  \n\n            except:\n                print('No weights loaded')    \n\n    return model","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:07.212885Z","iopub.execute_input":"2021-07-04T18:43:07.213374Z","iopub.status.idle":"2021-07-04T18:43:07.234487Z","shell.execute_reply.started":"2021-07-04T18:43:07.213324Z","shell.execute_reply":"2021-07-04T18:43:07.233054Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"base_model = compile_model(model=None, load_checkpoint=False, \n                           lr_scale_factor=500.0, label_smoothing=.1)","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:07.235964Z","iopub.execute_input":"2021-07-04T18:43:07.236386Z","iopub.status.idle":"2021-07-04T18:43:36.116280Z","shell.execute_reply.started":"2021-07-04T18:43:07.236342Z","shell.execute_reply":"2021-07-04T18:43:36.115031Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Test inference speed","metadata":{}},{"cell_type":"code","source":"def run_inference(model, dataset, return_lev_score=False, take_num=100, skip_set_num=0):\n    \"\"\" produces image_id, pred_string pairs \"\"\"\n    \n    # initialize containers\n    image_ids_list = []\n    generated_predictions_list = []\n    true_InChI_list = []\n    \n    # prepare dataset for parallel / distributed execution\n    if PARAMETERS.tpu():\n        if not take_num:  # use full dataset (~ 4 min on TPU)\n            dataset = PARAMETERS.strategy().experimental_distribute_dataset(dataset)\n        \n        else:  # use restricted dataset (useful for testing purposes)\n            dataset = PARAMETERS.strategy().experimental_distribute_dataset(\n                        dataset.skip(take_num * skip_set_num).take(take_num))\n            \n        # convert distributed ds to iterator\n        dataset = iter(dataset)\n\n    else:\n        if not take_num:  # use full dataset\n            pass\n        else:  # use restricted dataset (useful for testing purposes)\n            dataset = dataset.skip(take_num * skip_set_num).take(take_num).prefetch(tf.data.AUTOTUNE)\n        \n    # generate (image_id, token preds)\n    for val in dataset:\n\n        # get actual values and gather into single batch\n        true_InChI = val[3]\n        if PARAMETERS.tpu():\n            true_InChI = PARAMETERS.strategy().gather(true_InChI, axis=0)\n\n        # get corresponding image ids and gather into single batch\n        image_ids = val[2]\n        if PARAMETERS.tpu():\n            image_ids = PARAMETERS.strategy().gather(image_ids, axis=0)\n\n        # get predictions (as tokens)  and gather into single batch\n        if PARAMETERS.tpu():\n            generated_probs = PARAMETERS.strategy().run(model.generation_loop, args=[val[:2]])\n        else:\n            generated_probs = model.generation_loop(val[:2])\n        \n        if PARAMETERS.tpu():\n            generated_probs = PARAMETERS.strategy().gather(generated_probs, axis=0)\n        generated_probs = tf.squeeze(generated_probs)\n\n        # convert predictions to strings\n        generated_predictions = tf.argmax(generated_probs, axis=2)\n        generated_predictions = model.tokens_to_string(generated_predictions)\n\n        # decode bytestrings and update containers\n        image_ids_list.extend([x.decode() for x in image_ids.numpy().tolist()])\n        generated_predictions_list.extend([x.decode() for x in generated_predictions.numpy().tolist()])\n        true_InChI_list.extend([x.decode() for x in true_InChI.numpy().tolist()])\n\n    output = [image_ids_list, generated_predictions_list, true_InChI_list]\n\n    if return_lev_score:\n        \n        # compute scores\n        lev_score = [levenshtein(pred, orig) for (pred, orig)\n                      in zip(generated_predictions_list, true_InChI_list)]\n\n        # add to outputs\n        output.append(lev_score)\n\n    return output","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:36.117966Z","iopub.execute_input":"2021-07-04T18:43:36.118308Z","iopub.status.idle":"2021-07-04T18:43:36.135817Z","shell.execute_reply.started":"2021-07-04T18:43:36.118273Z","shell.execute_reply":"2021-07-04T18:43:36.134445Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# run_inference() not functional on Kaggle. Please use Colab instead.\n\"\"\"\nrun_inference(base_model, dataset=train_ds_int_index, return_lev_score=True, take_num=, skip_set_num=0)\n\"\"\"","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:36.137422Z","iopub.execute_input":"2021-07-04T18:43:36.137738Z","iopub.status.idle":"2021-07-04T18:43:36.163061Z","shell.execute_reply.started":"2021-07-04T18:43:36.137706Z","shell.execute_reply":"2021-07-04T18:43:36.161830Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"\n# test inference speed - time for 'take_num' batches\n%timeit run_inference(base_model, dataset=train_ds_int_index, return_lev_score=True, take_num=2, skip_set_num=0)\n\"\"\"","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:36.164451Z","iopub.execute_input":"2021-07-04T18:43:36.164744Z","iopub.status.idle":"2021-07-04T18:43:36.175590Z","shell.execute_reply.started":"2021-07-04T18:43:36.164716Z","shell.execute_reply":"2021-07-04T18:43:36.174376Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"markdown","source":"Train base model","metadata":{}},{"cell_type":"code","source":"# Dataset: Random Perturbations and One-Hot-Encoding\n# note: one-hot needed so we can use label smoothing in CrossEntropy Loss\ntf.keras.mixed_precision.set_global_policy('float32')  # temporarily removed mixed precision\n\ntrain_ds_prepared = train_ds_int_index\nvalid_ds_prepared = valid_ds_int_index\n\n# define transformations\n\"\"\"\nrotate = keras.layers.experimental.preprocessing.RandomRotation(\n            factor=(-0.5, 0.5), fill_mode='constant')\ncontrast = keras.layers.experimental.preprocessing.RandomContrast(factor=.1)\n\n# apply transformations\ntrain_ds_prepared = train_ds_prepared.map(lambda w, x, y, z: (rotate(w), x, y, z),\n                                           num_parallel_calls=tf.data.AUTOTUNE)\ntrain_ds_prepared = train_ds_prepared.map(lambda w, x: (contrast(w), x, y, z),\n                                          num_parallel_calls=tf.data.AUTOTUNE)\n\"\"\"\n\n# create one-hot encoded targets (allows label smoothing)\ndepth = base_model.vocab_size\none_hot = keras.layers.Lambda(lambda x: tf.one_hot(x, depth=depth)) \n\ntrain_ds_prepared = train_ds_prepared.map(lambda w, x, y, z: ((w, x), one_hot(x)),\n                                          num_parallel_calls=tf.data.AUTOTUNE)\\\n                                          .prefetch(tf.data.AUTOTUNE)\n\nvalid_ds_prepared = valid_ds_int_index.map(lambda w, x, y, z: ((w, x), one_hot(x)),\n                                          num_parallel_calls=tf.data.AUTOTUNE)\\\n                                          .prefetch(tf.data.AUTOTUNE)\n\n# re-enable to mixed precision\ntf.keras.mixed_precision.set_global_policy(PARAMETERS.mixed_precision())","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:43:36.177436Z","iopub.execute_input":"2021-07-04T18:43:36.178022Z","iopub.status.idle":"2021-07-04T18:43:36.259016Z","shell.execute_reply.started":"2021-07-04T18:43:36.177972Z","shell.execute_reply":"2021-07-04T18:43:36.257823Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"\nNote: validation results use inference step char generation. Training uses \nteacher-fed inputs.\n\"\"\"\n\n# Note: training accuracy based on teacher-training, not full generation loop\n\n# Note: label smoothing is helpful in training, but needs to be removed \n# in the final stages\n\n# recompile (optional)\n#base_model = compile_model(model=base_model, load_checkpoint=False, lr_scale_factor=500., label_smoothing=0.1)\n\n# Train base model (teacher-fed training, prediction-fed validation, no beam update)\nif not PARAMETERS.tpu():\n    steps_per_epoch = 512\n    validation_steps = 128\n    callbacks=[checkpoint, nan_stop, backup_checkpoint],# tensorboard]\n    validation_freq = 6\n\nelse:\n    steps_per_epoch = 12 * (8 * PARAMETERS.strategy().num_replicas_in_sync)\n    validation_steps = PARAMETERS.strategy().num_replicas_in_sync  \n    callbacks=[checkpoint, nan_stop]\n    validation_freq = 4\n\n\nepoch_multiple = 100\nepochs = epoch_multiple * int(1.8 * 1e6) // (steps_per_epoch * PARAMETERS.batch_size())\n    \n# (Optional: focused training.)\nbase_model.ImageFeaturesExtractor.trainable = True  \nbase_model.get_layer('EncoderAttention').trainable = True\nbase_model.get_layer('DecoderAttention').trainable = True\nbase_model.get_layer('DecoderHead').trainable = True\n\n# train\nhistory = base_model.fit(train_ds_prepared.repeat(),\n                         validation_data=valid_ds_prepared.repeat(),\n                         epochs=epochs,\n                         steps_per_epoch=steps_per_epoch,\n                         validation_freq=validation_freq, \n                         validation_steps=validation_steps, \n                         callbacks=callbacks,\n                         verbose=1)","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:47:42.831669Z","iopub.execute_input":"2021-07-04T18:47:42.832309Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"TPU-safe saving to local directory","metadata":{}},{"cell_type":"code","source":"base_model.save_weights(os.path.join(PARAMETERS.checkpoint_dir(), checkpoint_save_name, 'saved_model'), \n                        options=tf.saved_model.SaveOptions(experimental_io_device='/job:localhost'))","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:47:33.299900Z","iopub.status.idle":"2021-07-04T18:47:33.300370Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_base.load_weights(os.path.join(PARAMETERS.checkpoint_dir(), checkpoint_save_name, 'saved_model'), \n                        options=tf.saved_model.SaveOptions(experimental_io_device='/job:localhost'))","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:47:33.301489Z","iopub.status.idle":"2021-07-04T18:47:33.301917Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Train beam update model","metadata":{}},{"cell_type":"code","source":"\"\"\"\n# train beam model (prediction-fed training and inference, includes beam update mech)\n\nsteps_per_epoch = 150\nepochs = len(train_labels_df) // steps_per_epoch  # one full pass through the dataset\n\n# choose variables to train\nmodel_base.decoder_0.trainable = True\nmodel_base.decoder_1.trainable = True  # if multiple decoders enabled\n\n# compile\nmodel_beam.compile(optimizer=optimizer, \n              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=False),\n              metrics=['sparse_categorical_accuracy', EditDistanceMetric()])\n\n# train\nmodel_beam.fit(train_ds, epochs=30, steps_per_epoch=20, \n               validation_data=valid_ds, validation_steps=3, validation_freq=5,\n               callbacks=[checkpoint],#, tensorboard], \n               verbose=2, use_multiprocessing=True)\n\"\"\"","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:47:33.302752Z","iopub.status.idle":"2021-07-04T18:47:33.303296Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference\n\nHere we define function to conduct inference on the test set. Results are saved to \"submission.csv\".\n\nIntermediate results are saved at regular intervals to. This allows inference to be conducted in stages and is a safeguard in case of interruptions before the full set has been processed. ","metadata":{}},{"cell_type":"code","source":"def make_inference_progress(dataset, model, return_lev_score=True, save_freq=50, parameters=PARAMETERS):\n\n    batch_size = 1024\n    est_num_batches = 2*10e7 // batch_size\n    take_num = 100\n\n    #initialize new dataframe\n    predictions_df = pd.DataFrame(columns=['image_id', 'InChI', 'lev_score'])\n\n\n    for i in range(int(est_num_batches // take_num)):\n        try:\n\n             # get predictions\n            inference_outputs = run_inference(model, dataset, return_lev_score=True, \n                                              take_num=take_num, skip_set_num=i)\n            \n            im_id, pred, true_val, lev_score  = inference_outputs[:]\n\n            # add to dataframe\n            new_preds = pd.DataFrame({'image_id': im_id, 'InChI': pred, 'lev_score': lev_score})\n            predictions_df = predictions_df.append(new_preds)\n\n            # save to CSV\n            if i % save_freq == 0:\n                predictions_df = predictions_df.drop_duplicates(subset='image_id', keep='last')\n                predictions_df[['image_id', 'InChI']].to_csv(PARAMETERS.csv_save_dir() + 'submission.csv', index=False)\n                print(f'iteration {i}')\n\n        except:\n            print(f'completed at step {i}')\n            break\n\n    return predictions_df","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:47:33.304347Z","iopub.status.idle":"2021-07-04T18:47:33.304742Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Generate Predictions","metadata":{}},{"cell_type":"code","source":"predictions_df = make_inference_progress(predictions_df, save_freq=100, num_batches=1, starting_batch=0, parameters=PARAMETERS)\npredictions_df","metadata":{"execution":{"iopub.status.busy":"2021-07-04T18:47:33.305693Z","iopub.status.idle":"2021-07-04T18:47:33.306144Z"},"trusted":true},"execution_count":null,"outputs":[]}]}