{"cells":[{"metadata":{},"cell_type":"markdown","source":"

Bristol-Myers Squibb – Molecular Translation

\n

Image Captioning [TRAINING]

\n
CREATED BY: DARIEN SCHETTLER
\n\n---\n\n
\n"},{"metadata":{},"cell_type":"markdown","source":"

\n\n

\n\n

TABLE OF CONTENTS

\n\n---\n\n

0    IMPORTS

\n\n---\n\n

1    BACKGROUND INFORMATION

\n\n---\n\n

2    SETUP

\n\n---\n\n

3    HELPER FUNCTIONS

\n\n---\n\n

4    PREPARE THE DATASET

\n\n---\n\n

5    MODEL PREPARATION

\n\n---\n\n

6    DATASET CREATION

\n\n---\n\n

7    CUSTOM MODEL TRAINING

\n\n---\n\n

8    FINAL REMARKS

"},{"metadata":{},"cell_type":"markdown","source":"

0  IMPORTS    ⤒

"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# FOR KERAS UTILS PLOT\n!pip install -q pydot\n!pip install -q pydotplus\n!apt-get install -q graphviz\n!pip install -q git+https://github.com/qubvel/classification_models.git\n\nimport Levenshtein\n\nprint(\"\\n... IMPORTS STARTING ...\\n\")\nprint(\"\\n\\tVERSION INFORMATION\")\n\n# Machine Learning and Data Science Imports\nimport tensorflow as tf; print(f\"\\t\\t– TENSORFLOW VERSION: {tf.__version__}\");\nimport tensorflow_addons as tfa; print(f\"\\t\\t– TENSORFLOW ADDONS VERSION: {tfa.__version__}\");\nimport pandas as pd; pd.options.mode.chained_assignment = None;\nimport numpy as np; print(f\"\\t\\t– NUMPY VERSION: {np.__version__}\");\n\n# Imports for Alternative Models\nimport classification_models\nfrom classification_models.keras import Classifiers\n\n# Built In Imports\nfrom kaggle_datasets import KaggleDatasets\nfrom collections import Counter\nfrom datetime import datetime\nfrom glob import glob\nimport warnings\nimport requests\nimport imageio\nimport IPython\nimport urllib\nimport zipfile\nimport pickle\nimport random\nimport shutil\nimport string\nimport math\nimport time\nimport gzip\nimport ast\nimport io\nimport os\nimport gc\nimport re\n\n# Visualization Imports\nfrom matplotlib.colors import ListedColormap\nimport matplotlib.patches as patches\nimport plotly.graph_objects as go\nimport matplotlib.pyplot as plt\nfrom tqdm.notebook import tqdm; tqdm.pandas();\nimport plotly.express as px\nimport seaborn as sns\nfrom PIL import Image\nimport matplotlib; print(f\"\\t\\t– MATPLOTLIB VERSION: {matplotlib.__version__}\");\nimport plotly\nimport PIL\nimport cv2\n\n\nprint(\"\\n\\n... IMPORTS COMPLETE ...\\n\")\n\n\ngpus = tf.config.experimental.list_physical_devices('GPU')\nif gpus:\n try:\n # Currently, memory growth needs to be the same across GPUs\n for gpu in gpus:\n tf.config.experimental.set_memory_growth(gpu, True)\n logical_gpus = tf.config.experimental.list_logical_devices('GPU')\n print(len(gpus), \"... Physical GPUs,\", len(logical_gpus), \"Logical GPUs ...\\n\")\n except RuntimeError as e:\n # Memory growth must be set before GPUs have been initialized\n print(e)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

1  BACKGROUND INFORMATION    ⤒

\n\n
HIGH LEVEL TASK DESCRIPTION\n\n\n**Given an image, our goal is to generate a caption. In this case, that image is of a single molecule and the description/caption is the InChI string for that molecule.**\n\n---\n\n
\n\nMORE DETAIL ON IMAGE CAPTIONING\n\n\nDescription From a Tutorial I Used As Reference\n\n>Caption generation is a challenging artificial intelligence problem where a textual description must be generated for a given photograph.\n>\n>It requires both methods from computer vision to understand the content of the image and a language model from the field of natural language processing to turn the understanding of the image into words in the right order. Recently, deep learning methods have achieved state-of-the-art results on examples of this problem.\n>\n>Deep learning methods have demonstrated state-of-the-art results on caption generation problems. What is most impressive about these methods is a single end-to-end model can be defined to predict a caption, given a photo, instead of requiring sophisticated data preparation or a pipeline of specifically designed models.\n\n---\n\n
\n\nAPPROACH OVERVIEW\n\nIn this example, we will train a model on a relatively small amount of data to allow for us to see things end-to-end. We do this because the competition data is relatively large. We will also offer the ability to only infer on a smaller part of the InChI string to demonstrate capability."},{"metadata":{},"cell_type":"markdown","source":"

1.1 THE MODELS

\n\nThe overall model architecture we will use is similar to [**Show, Attend and Tell: Neural Image Caption Generation with Visual Attention**](https://arxiv.org/abs/1502.03044). Our pipeline is made up of 3 seperate components.\n\n
\n\nCOMPONENT 1 - THE FIXED ENCODER\n* This is the model that will take the image and encode it as a feature vector.\n* We will be using an off-the-shelf model from Tensorflow and we will NOT be training it (Using a model that is more fine-tuned for a task related to what we are doing would be a very good way to improve).\n* We will take the raw output of a headless backbone model and combine the smaller two axes together to give us a 2D vector output *(3D including batch)* **`(batch_size, attention_size, final_model_layer_channels_size)`**\n\nCOMPONENT 2 - THE TRAINABLE ENCODER\n* This is the model that will take the encoded feature vector from the **fixed encoder** and will pass it through a FC network to learn the associated task specific patterns.\n* We will be using a single FC layer and we WILL BE training it.\n* The output from this network will be the **IMAGE EMBEDDING** and it will have the following 3D Shape (including batch) **`(batch_size, attention_size, embedding_size)`**\n\nCOMPONENT 3 - THE TRAINABLE DECODER\n* This is the model that will take in multiple inputs and will pass them through a RNN (GRU/LSTM ... Transformer would be best probably) to learn how the image and trailing sequence informs on the desired prediction string.
\n      – **INPUT 1**: The feature vector encoded from the image of the molecule
\n      – **INPUT 2**: The [**hidden state**](https://ai.stackexchange.com/questions/16133/what-exactly-is-a-hidden-state-in-an-lstm-and-rnn) (i.e. what we've learned in the sequence so far)
\n      – **INPUT 3**: The previous word in the sequence (starts with the start token)

\n* We will be using a GRU layer for the Recurrent part of the network.\n* The actual network is very similar to [**Show, Attend and Tell Architecture**](https://arxiv.org/pdf/1502.03044.pdf)\n* The output from this network will be a vector the length of the vocabulary for every step (until a stop token is hit).\n\n"},{"metadata":{},"cell_type":"markdown","source":"

2  SETUP    ⤒

\n\nI put a lot of setup here that isn't required for this notebook. \n\nThe most important thing here is defining the list of tokens in such a way that when you pull the tokens from the InChI strings in order we do so correctly.\n\n---\n\n***NOTE:*** *The two token lists are given as **`ALL_POSSIBLE_CHARS`** which is for predicting on the entire InChI string and **`CHEM_FORM_CHARS`** which is the list of tokens if you only want to predict the chemical formula part of the InChI string.*"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Things that drastically change this notebook and are required\nLOAD_MODEL_FROM_FILE = True # Whether to load a pre/partially trained model from file\nTRAIN_AS_WE_GO = True # Whether to precompute the feature maps for the training dataset\nCHEM_ONLY = False # Whether to only try and predict the molecular formula\nN_TRAIN = 900_000 # Number of training examples for baseline\nN_VAL = 100_000 # Number of validation examples for baseline\n\n# Other Optional Information\nCKPT_PATH_LOAD = \"../input/bms-image-captioning-w-attention-train/model_training\" # Path to pre/partially trained model\nNPY_DIR_LOAD = \"\" # Path to previously computed feature embeddings for training data","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# DEEPLY UNNECESSARY TO INCLUDE THIS ARRAY OF ALL ELEMENTS\nALL_ELEMENTS = [\"H\", \" He\", \" Li\", \" Be\", \" B\", \" C\", \" N\", \" O\", \" F\", \" Ne\", \" Na\", \" Mg\", \" Al\", \" Si\", \" P\", \"S\", \" Cl\", \" Ar\", \" K\", \" Ca\", \" Sc\", \" Ti\", \" V\", \" Cr\", \" Mn\", \" Fe\", \" Co\", \" Ni\", \" Cu\", \" Zn\", \"Ga\", \" Ge\", \" As\", \" Se\", \" Br\", \" Kr\", \" Rb\", \" Sr\", \" Y\", \" Zr\", \" Nb\", \" Mo\", \" Tc\", \" Ru\", \"Rh\", \" Pd\", \" Ag\", \" Cd\", \" In\", \" Sn\", \" Sb\", \" Te\", \" I\", \" Xe\", \" Cs\", \" Ba\", \" La\", \" Ce\", \"Pr\", \" Nd\", \" Pm\", \" Sm\", \" Eu\", \" Gd\", \" Tb\", \" Dy\", \" Ho\", \" Er\", \" Tm\", \" Yb\", \" Lu\", \" Hf\",\"Ta\", \" W\", \" Re\", \" Os\", \" Ir\", \" Pt\", \" Au\", \" Hg\", \" Tl\", \" Pb\", \" Bi\", \" Po\", \" At\", \" Rn\", \"Fr\", \" Ra\", \" Ac\", \" Th\", \" Pa\", \" U\", \" Np\", \" Pu\", \" Am\", \" Cm\", \" Bk\", \" Cf\", \" Es\", \" Fm\", \"Md\", \" No\", \" Lr\", \" Rf\", \" Db\", \" Sg\", \" Bh\", \" Hs\", \" Mt\", \" Ds\", \" Rg\", \" Cn\", \" Uut\", \"Fl\", \" Uup\", \" Lv\", \" Uus\", \" Uuo\"]\nTRAIN_ELEMENTS = ['C', 'H', 'B', 'Br', 'Cl', 'F', 'I', 'N', 'O', 'P', 'S', 'Si']\nALL_POSSIBLE_CHARS = [\"\", \"InChI=1S/\", \"\", \"/c\", \"/h\", \"/m\", \"/t\", \"/b\", \"/s\", \"/i\"] +\\\n ['Si', 'Br', 'Cl', 'F', 'I', 'N', 'O', 'P', 'S', 'C', 'H', 'B', ] +\\\n [str(i) for i in range(167,-1,-1)] +\\\n [\"\\+\", \"\\(\", \"\\)\", \"\\-\", \",\", \"D\", \"T\"]\n\n# Whether or not to only try and caption the chemical formula part of the INCHI string\nCHEM_FORM_CHARS = [\"\", \"InChI=1S/\", \"\"] +\\\n ['Si', 'Br', 'Cl', 'F', 'I', 'N', 'O', 'P', 'S', 'C', 'H', 'B',] +\\\n [str(i) for i in range(167,-1,-1)]\n\n\n\n# Prefixes and How Common\n# -- ORDERING --> {c}{h/None}{b/None}{t/None}{m/None}{s/None}{i/None}{h/None}{t/None}{m/None}\nPREFIX_ORDERING = \"chbtmsihtm\"\n\n# Define the root and data directories\nROOT_DIR = \"/kaggle/input\"\nDATA_DIR = os.path.join(ROOT_DIR, \"bms-molecular-translation\")\n\nTRAIN_DIR = os.path.join(DATA_DIR, \"train\")\nTEST_DIR = os.path.join(DATA_DIR, \"test\")\n\nTRAIN_CSV_PATH = os.path.join(DATA_DIR, \"train_labels.csv\")\nSS_CSV_PATH = os.path.join(DATA_DIR, \"sample_submission.csv\")\n\ntrain_df = pd.read_csv(TRAIN_CSV_PATH)\ntrain_df[\"img_path\"] = train_df.image_id.apply(lambda x: os.path.join(TRAIN_DIR, x[0], x[1], x[2], x+\".png\"))\n\n# Let's simplify the task and just try and predict on the chemical structure\nif CHEM_ONLY:\n train_df[\"InChI_chem\"] = train_df.InChI.apply(lambda x: \"/\".join(x.split(\"/\", 2)[:2]))\n print(\"\\n... WE WILL BE TRYING TO PREDICT THE CHEMICAL FORMULA ONLY ...\\n\")\n\nprint(\"\\n... TRAIN DATAFRAME W/ PATHS ...\\n\")\ndisplay(train_df)\n\nss_df = pd.read_csv(SS_CSV_PATH)\nprint(\"\\n... SUBMISSION DATAFRAME ...\\n\")\ndisplay(ss_df)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

3  HELPER FUNCTIONS    ⤒

\n\nThese are mostly all necessary (maybe not **`flatten_l_o_l`**)... better doc strings will come in future versions"},{"metadata":{"trusted":true},"cell_type":"code","source":"def flatten_l_o_l(nested_list):\n \"\"\" Flatten a list of lists \"\"\"\n return [item for sublist in nested_list for item in sublist]\n\n\ndef tf_load_image(path, img_size=(224,224), tile_to_3_channel=True, invert=False):\n \"\"\" Load an image with the correct size and shape \"\"\"\n img = decode_img(tf.io.read_file(path), img_size, n_channels=1, invert=invert) \n if tile_to_3_channel:\n return tf.tile(img, tf.constant((1, 1, 3), dtype=tf.int32))\n else:\n return img\n\n \ndef pad_to_square(a, constant=255):\n \"\"\" Pad a tensor array `a` evenly until it is a square \"\"\"\n h_src = tf.shape(a)[0]\n w_src = tf.shape(a)[1] \n if w_src>h_src: # pad height\n n_to_add = w_src-h_src\n top_pad = n_to_add//2\n bottom_pad = n_to_add-top_pad\n a = tf.pad(a, [(top_pad, bottom_pad), (0, 0), (0, 0)], mode='constant', constant_values=constant)\n elif h_src>w_src: # pad width\n n_to_add = h_src-w_src\n left_pad = n_to_add//2\n right_pad = n_to_add-left_pad\n a = tf.pad(a, [(0, 0), (left_pad, right_pad), (0, 0)], mode='constant', constant_values=constant)\n else:\n pass\n return a\n \n \ndef decode_img(img, img_size=(224,224), n_channels=1, invert=False):\n \"\"\" Decode the image by utilizing TF ... pad to square ... and resize \"\"\"\n \n # convert the compressed string to a 3D uint8 tensor\n img = tf.image.decode_png(img, channels=n_channels)\n \n if invert:\n img = tf.ones_like(img, dtype=tf.uint8)*255-img\n constant_pad=0\n else:\n constant_pad=255\n \n # resize the image to the desired size\n img = pad_to_square(img, constant=constant_pad)\n img = tf.image.resize(img, img_size)\n return tf.cast(img, tf.uint8)\n\n\ndef tokens_to_str(caption_tokens, discard_padding=True):\n \"\"\" Should convert a string of token ids to an InChI string \"\"\"\n if discard_padding:\n return \"\".join([int_2_char_lex[x] for x in caption_tokens if x!=len(int_2_char_lex)])\n else:\n return \"\".join([int_2_char_lex[x] for x in caption_tokens])\n \n \ndef evaluate(image, from_np=False):\n \"\"\" TBD \"\"\"\n attention_plot = np.zeros((MAX_LEN, fixed_encoder.output_shape[1]))\n hidden = tf.zeros((1, RNN_UNITS), tf.float32)\n\n if not from_np:\n temp_input = tf.expand_dims(tf_load_image(image, img_size=INPUT_SHAPE[:-1]), 0)\n img_tensor_val = fixed_encoder(temp_input)\n else:\n img_tensor_val=image\n \n features = trainable_encoder(img_tensor_val)\n dec = tf.ones((1, 1), tf.uint8)\n result = [int_2_char_lex[1],]\n\n for i in range(MAX_LEN-1):\n predictions, hidden, attention_weights = dec_model([dec, hidden, features])\n attention_plot[i] = tf.reshape(attention_weights, (-1, )).numpy()\n predicted_id = tf.random.categorical(predictions, 1)[0][0].numpy()\n result.append(int_2_char_lex[predicted_id])\n if int_2_char_lex[predicted_id] == '':\n return result, attention_plot\n\n dec = tf.expand_dims([predicted_id], 0)\n\n attention_plot = attention_plot[:len(result), :]\n return result, attention_plot\n\n\ndef plot_attention(image, result, attention_plot):\n \"\"\" TBD \"\"\" \n temp_image = np.array(Image.open(image))\n\n fig = plt.figure(figsize=(18, 14))\n\n len_result = len(result)\n for l in range(len_result):\n temp_att = np.resize(attention_plot[l], (8, 8))\n ax = fig.add_subplot(len_result//2, len_result//2, l+1)\n ax.set_title(result[l])\n img = ax.imshow(temp_image)\n ax.imshow(temp_att, cmap='gray', alpha=0.4, extent=img.get_extent())\n\n plt.tight_layout()\n plt.show()\n \ndef test_random_image(style=\"full\"):\n \"\"\" TBD \"\"\" \n rid = np.random.randint(0, len(val_subset_df))\n path = val_subset_df[\"img_path\"][rid]\n \n if style==\"full\":\n real_caption = val_subset_df[\"InChI\"][rid][:-1]\n else:\n real_caption = val_subset_df[\"InChI_chem\"][rid][:-1]\n\n result, attention_plot = evaluate(path)\n result = ''.join(result[:-1])\n print (f\"\\n\\tReal Caption : {real_caption}\")\n print (f\"\\tPrediction Caption : {result}\")\n print(f\"\\tLevenshtein Distance : {Levenshtein.distance(real_caption, result)}\\n\")\n plot_attention(path, result, attention_plot)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

4  PREPARE THE DATASET    ⤒

\n\nIn this section we prepare a subset of the dataset for modelling"},{"metadata":{},"cell_type":"markdown","source":"

4.1 VISUALIZE SOME EXAMPLES

"},{"metadata":{"trusted":true},"cell_type":"code","source":"row_text = []\nplt.figure(figsize=(18,18))\nfor i,(row_idx,row) in enumerate(train_df.sample(9).iterrows()):\n row_text.append(f\"- {row.InChI}\")\n plt.subplot(3,3,i+1)\n plt.title(\"ROW #{}{}\".format(row_idx, '\\n'.join(row.InChI.replace(\"InChI=1S\", \"\").split(\"/\"))), fontweight='bold', fontsize=8)\n plt.imshow(tf_load_image(row.img_path))\n\nplt.tight_layout()\nplt.show()\n\nprint(\"\\n... CAPTIONING FOR IMAGES ABOVE ...\\n\")\nfor x in row_text:\n print(f\"{x}\")\nprint()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

4.2 CREATE SUBSET OF TRAINING DATA (ALSO CREATE VAL)

\n\nWe don't want to work on the whole dataset – **`2,424,186` IMAGES!!** – so we will subset the data. \n\n---\n\n***–– NOT DOING THIS SAMPLING METHOD IN THIS VERSION ––***
*        To try and capture all different types of data we will sample the data capturing InChI strings with all possible lengths.*
***–– NOT DOING THIS SAMPLING METHOD IN THIS VERSION ––***"},{"metadata":{"trusted":true},"cell_type":"code","source":"N_EX = len(train_df)\nprint(f\"# OF EXAMPLES : {N_EX}\")\nprint(f\"# OF TRAIN EXAMPLES : {N_TRAIN}\")\nprint(f\"# OF VALIDATION EXAMPLES : {N_VAL}\")\n\n# # Handles whether to sample in a more complicated fashion.\n# # This defaults to false if we are only doing the simpler demo notebook.\n# do_complicated = not CHEM_ONLY\n# if do_complicated:\n# train_df[\"inchi_len\"]= train_df.InChI.apply(lambda x: len(x))\n# # Get Ratios\n# len_ratios = {\n# k:v/len(train_df) \\\n# for k,v in train_df[\"inchi_len\"].value_counts().items()\n# }\n# # Use Approx Subset Size to Get The Number of Examples We Need For Each InChI Length\n# train_len_captures = {\n# k:int(np.ceil(v*APPROX_TRAIN_SUBSET)) \\\n# for k,v in len_ratios.items()\n# }\n# val_len_captures = {\n# k:int(np.ceil(v*APPROX_VAL_SUBSET)) \\\n# for k,v in len_ratios.items()\n# }\n# # Create the training dataframe we will use - subset_df\n# train_subset_df = pd.concat([\n# train_df[train_df.inchi_len==k][:v] \\\n# for k,v in train_len_captures.items()\n# ]).reset_index(drop=True).drop(columns=\"inchi_len\")\n# N_TRAIN = len(train_subset_df) # Also get the exact number of examples\n# val_subset_df = pd.concat([\n# train_df[train_df.inchi_len==k][v_train:(v_train+v)] \\\n# for (k,v),v_train in zip(val_len_captures.items(), train_len_captures.values())\n# ]).reset_index(drop=True).drop(columns=\"inchi_len\")\n# N_VAL = len(val_subset_df) # Also get the exact number of examples\n# else:\n\nval_subset_df = train_df[:N_VAL].reset_index(drop=True)\ntrain_subset_df = train_df[N_VAL:(N_TRAIN+N_VAL)].reset_index(drop=True)\n\nprint(\"\\n... TRAIN SUBSET ...\\n\")\ndisplay(train_subset_df)\n\nprint(\"\\n... VAL SUBSET ...\\n\")\ndisplay(val_subset_df)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

4.3 ADD STOP TOKENS TO THE CAPTIONS (INCHI STRINGS)

\n\nWe would normally also add **`START`** tokens... but because all the strings start with the string **`InChI=1S/`** we will consider this to be the **`START`** token."},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"\\n... TRAIN SUBSET ...\\n\")\ntrain_subset_df[\"InChI\"] = train_subset_df[\"InChI\"]+\"\"\nif CHEM_ONLY:\n train_subset_df[\"InChI_chem\"] = train_subset_df[\"InChI_chem\"]+\"\"\n for x in train_subset_df.InChI_chem[:5]: print(x)\nelse:\n for x in train_subset_df.InChI[:5]: print(x)\n \nprint(\"\\n... VAL SUBSET ...\\n\")\nval_subset_df[\"InChI\"] = val_subset_df[\"InChI\"]+\"\"\nif CHEM_ONLY:\n val_subset_df[\"InChI_chem\"] = val_subset_df[\"InChI_chem\"]+\"\"\n for x in val_subset_df.InChI_chem[:5]: print(x)\nelse:\n for x in val_subset_df.InChI[:5]: print(x)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

4.4 TOKENIZE THE CAPTION

\n\nWe now need to turn our strings into lists of tokens. We will use **`re.findall`** to accomplish this. We will also use **`re.findall`** to identify what the maximum length InChI string is within our data so we know what we will have to pad our sequences to."},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"\\n... Finding Max Length and Defining Certain Variables ...\\n\")\nif CHEM_ONLY:\n CHARS = CHEM_FORM_CHARS\n MAX_LEN = train_df.InChI_chem.progress_apply(lambda x: len(re.findall(\"|\".join(CHARS), x))).max()+1\nelse:\n CHARS = ALL_POSSIBLE_CHARS\n # MAX_LEN = train_df.InChI.progress_apply(lambda x: len(re.findall(\"|\".join(CHARS), x))).max()+1\n MAX_LEN = 277\n\n# Create tools to convert from token to string and back again.\nchar_2_int_lex = {c.strip(\"\\\\\"):i for i,c in enumerate(CHARS)}\nint_2_char_lex = {v:k for k,v in char_2_int_lex.items()}\n\n# I horribly use both of these variables and didn't want to check all my logic\n# but both of these are just the number of unique tokens in our vocabulary.\nVOCAB_LEN = VOCAB_SIZE = len(char_2_int_lex)\n\nprint(\"MAX FORMULA LENGTH : \", MAX_LEN)\nprint(\"LENGTH OF VOCAB : \", VOCAB_LEN)\nprint(\"CHAR LIST : \", CHARS, \"\\n\\n\")\n\nprint(\"\\n... Finding Tokenized String For All Examples ...\\n\")\n# Captions start as zeros because that is the padding token (by using zero we can mask it later)\ntrain_captions = np.zeros((len(train_subset_df), MAX_LEN,), dtype=np.uint8)\nval_captions = np.zeros((len(val_subset_df), MAX_LEN,), dtype=np.uint8)\n\n# Make the sparse, padded encodings for our captions \nfor captions, subset_df in zip([train_captions, val_captions], [train_subset_df, val_subset_df]):\n if CHEM_ONLY:\n for i, inchi in tqdm(enumerate(subset_df.InChI_chem.values), total=len(subset_df)):\n sparse_rep = [char_2_int_lex[c] for c in re.findall(\"|\".join(CHARS), inchi)]\n captions[i, :len(sparse_rep)] = sparse_rep\n else:\n for i, inchi in tqdm(enumerate(subset_df.InChI.values), total=len(subset_df)):\n sparse_rep = [char_2_int_lex[c] for c in re.findall(\"|\".join(CHARS), inchi)]\n captions[i, :len(sparse_rep)] = sparse_rep ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# COMPARE\nindices = [9,25000,899999]\n\nfor idx in indices:\n print(f\"\\nEXAMPLE SHOWING INDEX {idx}\\n\")\n print(\"RAW INCHI : \", train_subset_df[\"InChI\"][idx])\n print(\"TOKENS : \", \" \".join([int_2_char_lex[x] for x in train_captions[idx] if x!=0]))\n print(\"\\n\\n\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

5  MODEL PREPARATION    ⤒

"},{"metadata":{},"cell_type":"markdown","source":"

5.1 DEFINE ENCODER MODEL(S) AND PARAMETERS

\n\nWe define some parameters including the embedding dimensions and number of RNN units and input shape here.\n\nWe also create functions to return the two encoder models we previously discussed: **the fixed encoder** and **the trainable encoder**. We are currently using an **Inception ResNetV2** model, but any backbone would work."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Basic Training Parameters\nBATCH_SIZE = 128\nSHUFF_BUFF = 512\nEMBED_DIM = 256\nRNN_UNITS = 512\nAUTOTUNE = tf.data.AUTOTUNE\n\n# Reszize Input Images to...\nINPUT_SHAPE = (192,192,3)\n\n# Directory and Checkpoint Parameters\nCKPT_PATH = \"/kaggle/working/model_training\"\nos.makedirs(CKPT_PATH, exist_ok=True)\nif LOAD_MODEL_FROM_FILE:\n print(\"\\n... LOADING MODEL CHECKPOINT FROM PREVIOUSLY SESSION ...\")\n !cp {CKPT_PATH_LOAD}/* {CKPT_PATH}\nprint(f\"\\n... Checkpoints will be saved to the following directory:\\n\\t-->{CKPT_PATH}\\n\")\n\n# Numpy File Loading Parameters\nNPY_DIR = \"/kaggle/working/numpy_files\"\nos.makedirs(NPY_DIR, exist_ok=True)\nif LOAD_MODEL_FROM_FILE and TRAIN_AS_WE_GO:\n NPY_DIR = NPY_DIR_LOAD\n print(\"\\n... LOADING NPY FILES FROM PREVIOUSLY SESSION ...\\n\")\n \ndef get_encoder_fixed(bb, prep_fn, input_shape, freeze_up_to=\"\"):\n \"\"\" Return the fixed encoder - modular with tf.keras.applications \"\"\" \n if freeze_up_to!=\"\":\n for l in bb.layers:\n l.trainable=False\n if l.name==freeze_up_to: \n break\n \n inputs=tf.keras.layers.Input(shape=input_shape)\n prep_inputs = prep_fn(tf.cast(inputs, tf.float32))\n outputs = bb(prep_inputs)\n outputs = tf.keras.layers.Reshape((-1, outputs.shape[-1]))(outputs)\n _model = tf.keras.Model(inputs=[inputs], outputs=[outputs], name=\"fixed_encoder\")\n \n return _model\n\ndef get_encoder_trainable(input_shape, embedding_dim):\n \"\"\" Return the trainable encoder - Single FC Layer \"\"\" \n inputs=tf.keras.layers.Input(shape=input_shape)\n outputs = tf.keras.layers.Dense(embedding_dim, activation=\"relu\")(inputs)\n return tf.keras.Model(inputs=[inputs], outputs=[outputs], name=\"trainable_encoder\")\n\n# Cleanup just in case\ngc.collect(); tf.keras.backend.clear_session(); gc.collect(); gc.collect();\n\n# Variable backbone information\nmodel_factory, preprocessing_fn = Classifiers.get('resnet18')\nbackbone = model_factory(INPUT_SHAPE, weights='imagenet', include_top=False)\n\n# Two parts of the encoder model – fixed and trainable\nfixed_encoder = get_encoder_fixed(backbone, preprocessing_fn, INPUT_SHAPE, freeze_up_to=\"stage4_unit2_relu1\")\ntrainable_encoder = get_encoder_trainable(fixed_encoder.output_shape[1:], EMBED_DIM)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

5.2 CREATE NUMPY ENCODINGS OF IMAGES

\n\nSince we are not updating the weights in the fixed-encoder during model training. We can calculate the embeddings the model generates for each image **ONCE** and save them as numpy files ready to be used in the training loop. Alternatively we can create them as we go..."},{"metadata":{"trusted":true},"cell_type":"code","source":"if not TRAIN_AS_WE_GO:\n image_dataset = tf.data.Dataset.from_tensor_slices(train_subset_df[\"img_path\"].values)\n image_dataset = image_dataset.map(lambda x: (tf_load_image(x, img_size=INPUT_SHAPE[:-1]), x), num_parallel_calls=tf.data.AUTOTUNE) \\\n .batch(BATCH_SIZE) \\\n .prefetch(AUTOTUNE)\n\n for img, path in tqdm(image_dataset, total=len(train_subset_df)//BATCH_SIZE):\n batch_features = fixed_encoder(img)\n for bf, p in zip(batch_features, path):\n path_to_save_to = os.path.join(NPY_DIR, p.numpy().decode(\"utf-8\").rsplit(\"/\", 1)[1][:-4])\n np.save(path_to_save_to, bf.numpy())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

5.4 DEFINE DECODER MODEL AND PARAMETERS

\n\nThe model architecture is inspired by the [**Show, Attend and Tell** Paper](https://arxiv.org/pdf/1502.03044.pdf).\n\n---\n\n**The big takeaways are:**\n\n* Use attention ([**Bahdanau**](https://arxiv.org/abs/1409.0473) in this case)\n* Mask zeros in the embedding layer of the decoder\n* Hidden state is both an input and an output (let's us learn from the sequence as we traverse it)\n* We return the attention weights so we can visualize later"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Maybe like this\nBahdanauAttentionTFA = tfa.seq2seq.BahdanauAttention\n\n# Maybe like this though...\nclass BahdanauAttention(tf.keras.Model):\n \"\"\" TBD \"\"\"\n def __init__(self, units):\n super(BahdanauAttention, self).__init__()\n self.W1 = tf.keras.layers.Dense(units)\n self.W2 = tf.keras.layers.Dense(units)\n self.V = tf.keras.layers.Dense(1)\n\n def call(self, features, hidden):\n # features(CNN_encoder output) shape == (batch_size, 64, embedding_dim)\n\n # hidden shape == (batch_size, hidden_size)\n # hidden_with_time_axis shape == (batch_size, 1, hidden_size)\n hidden_with_time_axis = tf.expand_dims(hidden, 1)\n\n # attention_hidden_layer shape == (batch_size, 64, units)\n attention_hidden_layer = (tf.nn.tanh(self.W1(features) +\n self.W2(hidden_with_time_axis)))\n\n # score shape == (batch_size, 64, 1)\n # This gives you an unnormalized score for each image feature.\n score = self.V(attention_hidden_layer)\n\n # attention_weights shape == (batch_size, 64, 1)\n attention_weights = tf.nn.softmax(score, axis=1)\n\n # context_vector shape after sum == (batch_size, hidden_size)\n context_vector = attention_weights * features\n context_vector = tf.reduce_sum(context_vector, axis=1)\n\n return context_vector, attention_weights\n \ndef BahdanauAttn(units, style=\"custom\"):\n \"\"\" Get Attention Layer/Model \"\"\"\n if style is \"custom\":\n return BahdanauAttention(units)\n elif style is \"tfa\":\n return BahdanauAttentionTFA(units)\n else:\n raise ValueError(\"`style` argument must be one of ['custom', 'tfa']\")\n \ndef get_decoder(batch_size, n_units, enc_output_shape, vocab_size, embedding_dim):\n \"\"\" TBD \"\"\"\n input_x = tf.keras.layers.Input(shape=(1,))\n input_hidden = tf.keras.layers.Input(shape=(n_units,))\n input_features = tf.keras.layers.Input(shape=enc_output_shape)\n \n context_vector, attention_weights = BahdanauAttn(n_units)(input_features, input_hidden)\n x = tf.keras.layers.Embedding(vocab_size, embedding_dim, mask_zero=True)(input_x)\n x = tf.keras.layers.Concatenate(axis=-1)([tf.expand_dims(context_vector, 1), x])\n \n rnn_output, state = tf.keras.layers.GRU(n_units, return_sequences=True, return_state=True)(x)\n \n # shape == (batch_size, max_length, hidden_size)\n x = tf.keras.layers.Dense(n_units)(rnn_output)\n \n x = tf.keras.layers.Reshape((-1, x.shape[2]))(x)\n x = tf.keras.layers.Dense(vocab_size)(x)\n \n outputs = tf.keras.layers.Reshape((vocab_size,))(x)\n \n return tf.keras.Model(inputs=[input_x, input_hidden, input_features], outputs=[outputs, state, attention_weights])\n \ndec_model = get_decoder(BATCH_SIZE, RNN_UNITS, trainable_encoder.output_shape[1:], VOCAB_SIZE, EMBED_DIM)\nprint(\"\\n... DECODER INPUTS...\")\nfor i, x in enumerate(dec_model.inputs): print(f\"\\t--> INPUT {i+1} SHAPE = {x.shape}\")\n\nprint(\"\\n\\n\\n... DECODER MODEL SUMMARY...\\n\")\ndec_model.summary()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

5.5 DEFINE OPTIMIZER & LOSS

\n\nSince we are defining a custom training loop later in the tutorial, we will need a loss function that can take into consideration the padding token and mask out it's contribution (otherwise all we would do is predict )."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Training configuration\nEPOCHS = 3\nWARMUP_STEPS = 10000\nVERBOSE_FREQ = 100\nTRAIN_STEPS = int(np.ceil(N_TRAIN/BATCH_SIZE))\nTOTAL_STEPS = EPOCHS * TRAIN_STEPS\n\n# Sparse categorical cross entropy for sparse tokens\nloss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n from_logits=True, \n reduction='none',\n)\n\ndef loss_function(real, pred):\n \"\"\" TBD \"\"\"\n # Find the padding tokens\n mask = tf.math.logical_not(tf.math.equal(real, 0))\n \n # Calculate loss across everything using Sparse-CC\n loss_ = loss_object(real, pred)\n \n # Remove loss contribution from padding tokens\n mask = tf.cast(mask, dtype=loss_.dtype)\n loss_ *= mask\n \n # Take the average across the batches\n return tf.reduce_mean(loss_)\n\ndef lrfn(step, warmup_lr_start, lr_start, lr_final, decays, warmup_steps, total_steps):\n # exponential warmup\n if step < warmup_steps:\n warmup_factor = (step / warmup_steps) ** 2\n lr = warmup_lr_start + (lr_start - warmup_lr_start) * warmup_factor\n # staircase decay\n else:\n power = (step - warmup_steps) // ((total_steps - warmup_steps) / (decays + 1))\n decay_factor = ((lr_start / lr_final) ** (1 / decays)) ** power\n lr = lr_start / decay_factor\n\n return round(lr, 8)\n# plot the learning rate schedule\ndef plot_lr_schedule(lr_schedule, name):\n plt.figure(figsize=(15,8))\n plt.plot(lr_schedule)\n schedule_info = f'start: {lr_schedule[0]}, max: {max(lr_schedule)}, final: {lr_schedule[-1]}'\n plt.title(f'Step Learning Rate Schedule {name}, {schedule_info}', size=16)\n plt.grid()\n plt.show()\n\n# Learning rate for encoder\nLR_SCHEDULE = [\n lrfn(\n step=step, \n warmup_lr_start=1e-6, \n lr_start=15e-4, \n lr_final=35e-5, \n decays=EPOCHS, \n warmup_steps=WARMUP_STEPS, \n total_steps=TOTAL_STEPS) \\\n for step in range(TOTAL_STEPS)\n]\n\n# custom learning rate scheduler\nclass LRReduce():\n def __init__(self, optimizer, lr_schedule):\n self.opt = optimizer\n self.lr_schedule = lr_schedule\n \n # assign initial learning rate\n self.lr = lr_schedule[0]\n \n self.opt.learning_rate.assign(self.lr)\n \n def step(self, step):\n self.lr = self.lr_schedule[step]\n \n # assign learning rate to optimizer\n self.opt.learning_rate.assign(self.lr)\n \n def get_counter(self):\n return self.c\n \n def get_lr(self):\n return self.lr\n\nOPTIMIZER = tf.keras.optimizers.Adam()\nLR_reduce = LRReduce(OPTIMIZER, LR_SCHEDULE)\nplot_lr_schedule(LR_SCHEDULE, 'Encoder')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

6  DATASET CREATION    ⤒

\n\nWe now have all the component parts of the dataset (and model parts) ready to go. We will now make a **`tf.data.Dataset`** for training and validation."},{"metadata":{"trusted":true},"cell_type":"code","source":"def np_map_func(img_name, caption):\n \"\"\" Load numpy from file \"\"\"\n img_tensor = np.load(os.path.join(NPY_DIR, img_name.decode('utf-8').rsplit(\"/\", 1)[1][:-4]+'.npy'))\n return img_tensor, caption\n\ndef train_as_go_preprocessing_fn(img_name, caption):\n img_tensor = tf_load_image(img_name, img_size=(192,192), tile_to_3_channel=True, invert=True)\n return img_tensor, caption\n \n# Train Dataset\ntrain_features_ds = tf.data.Dataset.from_tensor_slices(train_subset_df[\"img_path\"].values)\ntrain_captions_ds = tf.data.Dataset.from_tensor_slices(train_captions)\ntrain_ds = tf.data.Dataset.zip((train_features_ds, train_captions_ds))\n\n# Validation Dataset\nval_features_ds = tf.data.Dataset.from_tensor_slices(val_subset_df[\"img_path\"].values)\nval_captions_ds = tf.data.Dataset.from_tensor_slices(val_captions)\nval_ds = tf.data.Dataset.zip((val_features_ds, val_captions_ds))\n\nprint(\"\\n... RAW EXAMPLE ...\\n\")\nfor x,y in val_ds.take(1): print(f\"{x}\\n{y}\")\n \nprint(\"\\n... PARSED EXAMPLE ...\\n\")\nfor x,y in val_ds.take(1): print(f\"{x.numpy().decode()}\\n{tokens_to_str(y.numpy())}\")\n \nif TRAIN_AS_WE_GO:\n train_ds = train_ds.map(train_as_go_preprocessing_fn, num_parallel_calls=AUTOTUNE) \\\n .shuffle(SHUFF_BUFF) \\\n .batch(BATCH_SIZE) \\\n .prefetch(AUTOTUNE)\n val_ds = val_ds.map(train_as_go_preprocessing_fn, num_parallel_calls=AUTOTUNE) \\\n .batch(BATCH_SIZE) \\\n .prefetch(AUTOTUNE)\nelse:\n # Consume numpy files using tf.numpy_function\n train_ds = train_ds.map(lambda item1, item2: tf.numpy_function(\n np_map_func, [item1, item2], [tf.float32, tf.uint8]),\n num_parallel_calls=tf.data.AUTOTUNE) \\\n .shuffle(SHUFF_BUFF) \\\n .batch(BATCH_SIZE) \\\n .prefetch(AUTOTUNE)\n val_ds = val_ds.map(lambda item1, item2: tf.numpy_function(\n np_map_func, [item1, item2], [tf.float32, tf.uint8]),\n num_parallel_calls=tf.data.AUTOTUNE) \\\n .batch(BATCH_SIZE) \\\n .prefetch(AUTOTUNE)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

7  CUSTOM MODEL TRAINING    ⤒

\n\nWe will now define the component parts of our training loop to allow us to train our model"},{"metadata":{},"cell_type":"markdown","source":"

7.1 CHECKPOINTING

"},{"metadata":{"trusted":true},"cell_type":"code","source":"ckpt = tf.train.Checkpoint(fixed_encoder=fixed_encoder,\n trainable_encoder=trainable_encoder,\n decoder=dec_model,\n optimizer=OPTIMIZER)\nckpt_manager = tf.train.CheckpointManager(ckpt, CKPT_PATH, max_to_keep=5)\nstart_epoch = 0\nif ckpt_manager.latest_checkpoint:\n print(\"\\n... RESTORING FROM PREVIOUS CHECKPOINT ...\\n\")\n start_epoch = int(ckpt_manager.latest_checkpoint.split('-')[-1])\n\n # restoring the latest checkpoint in checkpoint_path\n ckpt.restore(ckpt_manager.latest_checkpoint)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

7.2 A SINGLE TRAIN STEP

"},{"metadata":{"trusted":true},"cell_type":"code","source":"@tf.function\ndef train_step_wo_encoding(img_tensor, target):\n \"\"\" This is ONE step in training \"\"\"\n # Step 1:\n # --> Initialize the loss for the batch to 0\n loss = 0.0\n\n # Step 2:\n # --> Initialize the hidden state (INPUT 1)\n # --> Initialize the decoder input (previous word) by setting it to\n # one. Therefore we get a vector of ones with the shape (batch_size, 1) (INPUT 2)\n _hidden = tf.zeros((target.shape[0], RNN_UNITS), tf.float32)\n _dec = tf.ones((target.shape[0], 1), tf.uint8)\n\n # Step 3-6 - Use gradient to record operations for automatic differentiation.\n with tf.GradientTape() as tape:\n # Step 3:\n # --> Calculate the feature encoding for the batch of images (INPUT 3)\n # - NOTE: We only do this once per batch\n _features = trainable_encoder(img_tensor, training=True)\n \n # Step 4-6 - Iterate over the tokens in the ground truth InChI caption \n # - Tokenized caption is padded and contains end token\n # - We start from i==1 as we initialized to i==0 (start token)\n for i in tf.range(1, target.shape[1]):\n # Step 4:\n # --> Pass the inputs to the decoder model\n # --> Capture the predicted next token and the updated hidden state\n # - NOTE: We will do this N times per batch (N is the max length of a tokenized caption)\n predictions, _hidden, _ = dec_model([_dec, _hidden, _features], training=True)\n \n # Step 5:\n # --> Calculate and add the loss for this token's prediction to the running loss total\n # - NOTE: We will do this N times per batch (N is the max length of a tokenized caption)\n loss += loss_function(target[:, i], predictions)\n\n # Step 6:\n # --> Teacher Forcing (use the ground truth from a prior time step as input)\n # --> i.e. We update the word to be passed to the decoder next...\n # - NOTE: We will do this N times per batch (N is the max length of a tokenized caption)\n _dec = tf.expand_dims(target[:, i], 1)\n \n # Step 7:\n # --> Calculate the average loss across the tokenized caption. This is the batchwise loss.\n total_loss = (loss/tf.cast(target.shape[1], tf.float32))\n\n # Step 8:\n # --> Capture which variables within the models we want to update on (conduct auto-diff on)\n trainable_variables = trainable_encoder.trainable_variables + dec_model.trainable_variables\n\n # Step 9:\n # --> Conduct automatic differentiation w.r.t. each variable in training_variables\n gradients = tape.gradient(loss, trainable_variables)\n\n # Step 10:\n # --> Update the optimizer by applying the gradients\n OPTIMIZER.apply_gradients(zip(gradients, trainable_variables))\n\n # Step 11:\n # --> Return loss values\n return loss, total_loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@tf.function\ndef train_step_w_encoding(img_tensor, target):\n \"\"\" This is ONE step in training \"\"\"\n # Step 1:\n # --> Initialize the loss for the batch to 0\n loss = 0.0\n\n # Step 2:\n # --> Initialize the hidden state (INPUT 1)\n # --> Initialize the decoder input (previous word) by setting it to\n # one. Therefore we get a vector of ones with the shape (batch_size, 1) (INPUT 2)\n _hidden = tf.zeros((target.shape[0], RNN_UNITS), tf.float32)\n _dec = tf.ones((target.shape[0], 1), tf.uint8)\n\n # Step 3-6 - Use gradient to record operations for automatic differentiation.\n with tf.GradientTape() as tape:\n # Step 3:\n # --> Calculate the feature encoding for the batch of images (INPUT 3)\n # - NOTE: We only do this once per batch\n img_embedding = fixed_encoder(img_tensor)\n _features = trainable_encoder(img_embedding, training=True)\n \n # Step 4-6 - Iterate over the tokens in the ground truth InChI caption \n # - Tokenized caption is padded and contains end token\n # - We start from i==1 as we initialized to i==0 (start token)\n for i in tf.range(1, target.shape[1]):\n # Step 4:\n # --> Pass the inputs to the decoder model\n # --> Capture the predicted next token and the updated hidden state\n # - NOTE: We will do this N times per batch (N is the max length of a tokenized caption)\n predictions, _hidden, _ = dec_model([_dec, _hidden, _features], training=True)\n \n # Step 5:\n # --> Calculate and add the loss for this token's prediction to the running loss total\n # - NOTE: We will do this N times per batch (N is the max length of a tokenized caption)\n loss += loss_function(target[:, i], predictions)\n\n # Step 6:\n # --> Teacher Forcing (use the ground truth from a prior time step as input)\n # --> i.e. We update the word to be passed to the decoder next...\n # - NOTE: We will do this N times per batch (N is the max length of a tokenized caption)\n _dec = tf.expand_dims(target[:, i], 1)\n \n # Step 7:\n # --> Calculate the average loss across the tokenized caption. This is the batchwise loss.\n total_loss = (loss/tf.cast(target.shape[1], tf.float32))\n\n # Step 8:\n # --> Capture which variables within the models we want to update on (conduct auto-diff on)\n trainable_variables = fixed_encoder.trainable_variables + trainable_encoder.trainable_variables + dec_model.trainable_variables\n\n # Step 9:\n # --> Conduct automatic differentiation w.r.t. each variable in training_variables\n gradients = tape.gradient(loss, trainable_variables)\n\n # Step 10:\n # --> Update the optimizer by applying the gradients\n OPTIMIZER.apply_gradients(zip(gradients, trainable_variables))\n\n # Step 11:\n # --> Return loss values\n return loss, total_loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def evaluate(image, from_np=False):\n \"\"\" TBD \"\"\"\n attention_plot = np.zeros((MAX_LEN, fixed_encoder.output_shape[1]))\n hidden = tf.zeros((1, RNN_UNITS), tf.float32)\n\n if not from_np:\n temp_input = tf.expand_dims(tf_load_image(image, img_size=INPUT_SHAPE[:-1]), 0)\n img_tensor_val = fixed_encoder(temp_input)\n else:\n img_tensor_val=image\n \n features = trainable_encoder(img_tensor_val)\n dec = tf.ones((1, 1), tf.uint8)\n result = [int_2_char_lex[1],]\n\n for i in range(MAX_LEN-1):\n predictions, hidden, attention_weights = dec_model([dec, hidden, features])\n attention_plot[i] = tf.reshape(attention_weights, (-1, )).numpy()\n predicted_id = tf.random.categorical(predictions, 1)[0][0].numpy()\n result.append(int_2_char_lex[predicted_id])\n if int_2_char_lex[predicted_id] == '':\n return result, attention_plot\n\n dec = tf.expand_dims([predicted_id], 0)\n\n attention_plot = attention_plot[:len(result), :]\n return result, attention_plot\n\n\ndef plot_attention(image, result, attention_plot, limit_n=9):\n \"\"\" TBD \"\"\" \n temp_image = tf_load_image(image)\n\n fig = plt.figure(figsize=(20, 20))\n\n len_result = len(result)\n \n if len_result>limit_n:\n print(f\"\\n...CAPPING ATTENTION MAP VISUALIZATION TO {limit_n} IMAGES AS DISPLAYING {len_result} IMAGES IS UNWIELDY...\\n\")\n len_result=limit_n\n attention_plot = attention_plot[:limit_n, :]\n sp_w = min(len_result//2, 3)\n sp_h = int(np.ceil(len_result/sp_w))\n for l in range(len_result):\n temp_att = np.resize(attention_plot[l], (8, 8))\n ax = fig.add_subplot(sp_h, sp_w, l+1)\n ax.set_title(result[l])\n img = ax.imshow(temp_image, cmap=\"gray\")\n ax.imshow(temp_att, cmap='gray', alpha=0.4, extent=img.get_extent())\n\n plt.tight_layout()\n plt.show()\n \ndef test_random_image(style=\"full\"):\n \"\"\" TBD \"\"\" \n rid = np.random.randint(0, len(val_subset_df))\n path = val_subset_df[\"img_path\"][rid]\n \n if style==\"full\":\n real_caption = val_subset_df[\"InChI\"][rid].replace(\"\", \"\")\n else:\n real_caption = val_subset_df[\"InChI_chem\"][rid].replace(\"\", \"\")\n\n result, attention_plot = evaluate(path)\n result_str = ''.join(result[:-1])\n print (f\"\\n\\tReal Caption : {real_caption}\")\n print (f\"\\tPrediction Caption : {result_str}\")\n print(f\"\\tLevenshtein Distance : {Levenshtein.distance(real_caption, result_str)}\\n\")\n plot_attention(path, result, attention_plot)\n \ndef lev_on_validation(ds, from_np=False, from_distribution=False):\n lev_distances = []\n for inp, captions in tqdm(ds):\n # Decoder Inputs\n dec = tf.ones((BATCH_SIZE, 1), tf.uint8)\n hidden = tf.zeros((BATCH_SIZE, RNN_UNITS), tf.float32)\n \n if not from_np:\n inp = fixed_encoder(inp)\n features = trainable_encoder(inp)\n\n # Initialize prediction array\n pred_captions = [[int_2_char_lex[1],] for b_sub in range(BATCH_SIZE)]\n\n # Capture ground truth captions\n gt_captions = [\"\".join([int_2_char_lex[c] for c in cap if c not in [0, 2]]) for cap in captions.numpy()]\n\n for i in range(MAX_LEN-1):\n predictions, hidden, _ = dec_model([dec, hidden, features])\n if from_distribution:\n predicted_ids = [tf.random.categorical(tf.expand_dims(preds,axis=0), 1)[0][0].numpy() for preds in predictions]\n else:\n predicted_ids = [tf.argmax(preds).numpy() for preds in predictions]\n for b_sub in range(BATCH_SIZE):\n pred_captions[b_sub].append(int_2_char_lex[predicted_ids[b_sub]])\n dec = tf.expand_dims(predicted_ids, -1)\n\n pred_captions = [\"\".join([c for c in caption_list if c != \"\"]) for caption_list in pred_captions]\n lev_distances+=[Levenshtein.distance(gt_c, pred_c) for gt_c, pred_c in zip(gt_captions, pred_captions)]\n \n # return lev_distances, np.array(lev_distances).mean()\n return np.array(lev_distances).mean()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"lev_distances = []\nfor inp, captions in tqdm(train_ds.take(1)):\n # Decoder Inputs\n dec = tf.ones((BATCH_SIZE, 1), tf.uint8)\n hidden = tf.zeros((BATCH_SIZE, RNN_UNITS), tf.float32)\n\n inp = fixed_encoder(inp)\n features = trainable_encoder(inp)\n\n # Initialize prediction array\n pred_captions = [[int_2_char_lex[1],] for b_sub in range(BATCH_SIZE)]\n\n # Capture ground truth captions\n gt_captions = [\"\".join([int_2_char_lex[c] for c in cap if c not in [0, 2]]) for cap in captions.numpy()]\n\n for i in range(MAX_LEN-1):\n predictions, hidden, _ = dec_model([dec, hidden, features])\n predicted_ids = [tf.random.categorical(tf.expand_dims(preds,axis=0), 1)[0][0].numpy() for preds in predictions]\n for b_sub in range(BATCH_SIZE):\n pred_captions[b_sub].append(int_2_char_lex[predicted_ids[b_sub]])\n dec = tf.expand_dims(predicted_ids, -1)\n\n pred_captions = [\"\".join([c for c in caption_list if c != \"\"]) for caption_list in pred_captions]\n lev_distances+=[Levenshtein.distance(gt_c, pred_c) for gt_c, pred_c in zip(gt_captions, pred_captions)]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dec = tf.ones((BATCH_SIZE, 1), tf.uint8)\nhidden = tf.zeros((BATCH_SIZE, RNN_UNITS), tf.float32)\nfor i in tqdm(range(MAX_LEN-1)):\n predictions, hidden, _ = dec_model([dec, hidden, features])\n predicted_ids = [tf.argmax(preds).numpy() for preds in predictions]\n for b_sub in range(BATCH_SIZE):\n pred_captions[b_sub].append(int_2_char_lex[predicted_ids[b_sub]])\n dec = tf.expand_dims(predicted_ids, -1)\nprint(\"\".join(pred_captions[0]))\ngt_caption","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

7.3 CONDUCT TRAINING

"},{"metadata":{"trusted":true},"cell_type":"code","source":"if CHEM_ONLY:\n style=\"chem\"\nelse:\n style=\"full\"\n \nloss_plot = []\nstep_total = 0\nfor epoch in range(start_epoch, EPOCHS):\n start = time.time()\n total_loss = 0.0\n \n print(f\"\\n... ATTEMPT ON RANDOM VALIDATION IMAGE PRIOR TO EPOCH #{epoch+1} ...\\n\")\n try:\n test_random_image(style=style)\n except:\n print(\"\\n... PLOTTING FAILED ...\\n\")\n\n for (batch, (img_tensor, target)) in enumerate(train_ds):\n step_total += 1\n batch_loss, t_loss = train_step_w_encoding(img_tensor, target)\n total_loss += t_loss\n\n if batch % 40 == 0:\n print (f\"\\n\\tEpoch: {epoch+1}\\n\\tBatch: {batch}\\n\\tBatch Loss: {batch_loss.numpy()/int(target.shape[1]):.4f}\\n\\tTotal Loss: {total_loss.numpy()/int(batch+1)}\")\n gc.collect(); gc.collect();\n \n LR_reduce.step(step_total)\n \n if np.isnan(total_loss):\n print(\"\\n\\n\\n\\n\\n\\n\\n\\n\\n\\nNANANANANANANANANAN\\nNANANANANANANANANAN\\nNANANANANANANANANAN\\n\\n\\n\\n\\n\\n\\n\\n\\n\\n\")\n break\n \n # storing the epoch end loss value to plot later\n loss_plot.append(total_loss/TRAIN_STEPS)\n \n ckpt_manager.save()\n \n print(f\"\\n\\nEpoch: {epoch+1}\\nEpoch Loss: {total_loss/TRAIN_STEPS:.6f}\")\n print(f\"Lev Distance on Subset of Data: {lev_on_validation(val_ds.take(10))}\")\n print(f\"Time taken for 1 epoch {time.time()-start} sec\\n\\n\")\n \nprint(f\"Lev Distance on All Validation Data: {lev_on_validation(val_ds)}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

8  FINAL REMARKS & NEXT STEPS    ⤒

\n\nNEXT STEPS

\n     - Create a better **fixed encoder**
\n     - Add additional layers to **trainable encoder**
\n     - Improve the **decoder** (more layers, different types, Transformer, etc.)
\n     - Preprocess the images and remove noise in the background.
\n     - Invert so that white background becomes black (0s) and foreground becomes white (larger values for more distinct areas/lines).
\n     - Train on more than just **`~1.5-2%`** of the data
\n     - Attempt to predict more than just the chemical formula
\n     - Lots more...\n\n
\n\n---\n\n
\n\nFINAL REMARKS

\n\n\n \n
This notebook is a learning tool for me that I wanted to share with everyone.
\n\n\n \n
If there are ways I can improve or things you think I can do better please let me know.
\n\n
Thanks for taking the time and reading this far !!

\n\n

"}],"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":4,"nbformat_minor":4}