{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"%%capture\n!pip install timm\n!pip install python-Levenshtein","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport re\nimport time\n\nimport pandas as pd\nimport numpy as np\n\nimport timm\n\nimport cv2\n\nfrom tqdm import tqdm\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.nn.utils.rnn import pad_sequence\nfrom torch.nn.utils.rnn import pack_padded_sequence\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.backends.cudnn as cudnn\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nfrom sklearn.model_selection import train_test_split\n\nfrom Levenshtein import distance as levenshtein_distance","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"BASE_DIR = '../input/bms-molecular-translation/train'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"label_df = pd.read_csv(\"../input/bms-molecular-translation/train_labels.csv\")\nlabel_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Remove \"InchI=1S/\" as it's common for all notations\nlabel_df[\"InChI\"] = label_df[\"InChI\"].str[9: ]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nformula_regex = re.compile(r\"(.*?)/c\")\natom_conn_regex = re.compile(r\"/c(.*?)(/h|$)\")\nhydrogen_regex = re.compile(r\"/h(.*)\")\n\ndef get_formulae(inChI):\n    try:\n        return formula_regex.search(inChI).group(1)\n    except:\n        return None\n\ndef get_atom_conn(inChI):\n    try:\n        return atom_conn_regex.search(inChI).group(1)\n    except:\n        return None\n\ndef get_hydrogen(inChI):\n    try:\n        return hydrogen_regex.search(inChI).group(1)\n    except:\n        return None\n\nlabel_df[\"formula\"] = label_df[\"InChI\"].map(get_formulae)\nlabel_df[\"atom_conn\"] = label_df[\"InChI\"].map(get_atom_conn)\nlabel_df[\"hydrogen\"] = label_df[\"InChI\"].map(get_hydrogen)\n\nlabel_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"label_df.isna().sum()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"There are a few training samples having no hydrogen atom information in InChI notation. Dropping those."},{"metadata":{"trusted":true},"cell_type":"code","source":"label_df = label_df.dropna(axis=0)\nlabel_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Currently using only 90k+10k images\nlabel_df = label_df.iloc[:100000]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Train on 90k, validate on 10k\ntrain_df, test_df = train_test_split(label_df, test_size=0.1, shuffle=False)\ntrain_df = train_df.reset_index(drop=True)\ntest_df = test_df.reset_index(drop=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df.shape, test_df.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Build Vocabulary"},{"metadata":{},"cell_type":"markdown","source":"There are only limited set of characters in the InChI notation. Refer this for more details: https://www.kaggle.com/c/bms-molecular-translation/discussion/223471. So I'll build a character based vocabulary class which will be used for target preprocessing."},{"metadata":{"trusted":true},"cell_type":"code","source":"class Vocabulary:\n    def __init__(self, freq_threshold=2, reverse=False):\n        self.itos = {0: \"<pad>\", 1: \"<sos>\", 2: \"<eos>\", 3: \"<unk>\"}\n        self.stoi = {\"<pad>\": 0, \"<sos>\": 1, \"<eos>\": 2, \"<unk>\": 3}\n        self.freq_threshold = freq_threshold\n        self.reverse = reverse\n        self.tokenizer = self._tokenizer\n\n    def __len__(self):\n        return len(self.itos)\n    \n    def _tokenizer(self, text):\n        return (char for char in text)\n\n    def tokenize(self, text):\n        if self.reverse:\n            return [token for token in self.tokenizer(text)][::-1]\n        else:\n            return [token for token in self.tokenizer(text)]\n\n    def build_vocabulary(self, sentence_list):\n        \"\"\"Basically builds a frequency map for all possible characters.\"\"\"\n        frequencies = {}\n        idx = len(self.itos)\n\n        for sentence in sentence_list:\n            # Preprocess the InChI.\n            for char in self.tokenize(sentence):\n                if char in frequencies:\n                    frequencies[char] += 1\n                else:\n                    frequencies[char] = 1\n\n                if frequencies[char] == self.freq_threshold:\n                    self.stoi[char] = idx\n                    self.itos[idx] = char\n                    idx += 1\n\n    def numericalize(self, text):\n        \"\"\"Convert characters to numbers.\"\"\"\n        tokenized_text = self.tokenize(text)\n\n        return [\n            self.stoi[token] if token in self.stoi else self.stoi[\"<unk>\"]\n            for token in tokenized_text\n        ]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n# Build vocab using training data\nfreq_threshold = 2\nvocab_formula = Vocabulary(freq_threshold=freq_threshold, reverse=False)\nvocab_atom_conn = Vocabulary(freq_threshold=freq_threshold, reverse=False)\nvocab_hydrogen = Vocabulary(freq_threshold=freq_threshold, reverse=False)\n\n# build vocab\nvocab_formula.build_vocabulary(train_df['formula'].to_list())\nvocab_atom_conn.build_vocabulary(train_df['atom_conn'].to_list())\nvocab_hydrogen.build_vocabulary(train_df['hydrogen'].to_list())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"IMG_SIZE = 128","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Augmentations"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_augs():\n    return A.Compose(\n        [\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.ShiftScaleRotate(p=0.5),\n            A.GaussNoise(p=0.5),\n            A.Cutout(num_holes=8, max_h_size=64, max_w_size=64, fill_value=255, p=0.5),\n            ToTensorV2(p=1.0),\n        ], \n        p=1.0\n    )\n\ndef get_valid_augs():\n    return A.Compose(\n        [\n            ToTensorV2(p=1.0),\n        ], \n        p=1.0\n    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class BMSDataset(Dataset):\n    def __init__(self, df, augs, vocab_formula, vocab_atom_conn, vocab_hydrogen):\n        super().__init__()\n        self.df = df\n        self.augs = augs\n        self.vocab_formula = vocab_formula\n        self.vocab_atom_conn = vocab_atom_conn\n        self.vocab_hydrogen = vocab_hydrogen\n\n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        # Read the image\n        img_id = self.df.iloc[idx]['image_id']\n        label = self.df.iloc[idx]['InChI']\n        \n        # Split formula, atom connection and hydrogen sections\n        label_formula = self.df.iloc[idx]['formula']\n        label_atom_conn = self.df.iloc[idx]['atom_conn']\n        label_hydrogen = self.df.iloc[idx]['hydrogen']\n        \n        label_formula_len = len(label_formula) + 2  # (2 for <sos> and <eos>)\n        label_atom_conn_len = len(label_atom_conn) + 2\n        label_hydrogen_len = len(label_hydrogen) + 2\n        \n        img_path = os.path.join(BASE_DIR, img_id[0], img_id[1], img_id[2], f'{img_id}.png')\n        \n        img = self._load_from_file(img_path)\n        \n        # Apply image augmentations if available\n        if self.augs:\n            img = self.augs(image=img)[\"image\"]\n        \n        # Convert label to numbers\n        label_formula = self._get_numericalized(label_formula, self.vocab_formula)\n        label_atom_conn = self._get_numericalized(label_atom_conn, self.vocab_atom_conn)\n        label_hydrogen = self._get_numericalized(label_hydrogen, self.vocab_hydrogen)\n\n        return img, torch.tensor(label_formula), torch.tensor(label_atom_conn), torch.tensor(label_hydrogen), torch.tensor(label_formula_len), torch.tensor(label_atom_conn_len), torch.tensor(label_hydrogen_len)\n    \n    def _get_numericalized(self, sentence, vocab):\n        \"\"\"Numericalize given text using prebuilt vocab.\"\"\"\n        numericalized = [vocab.stoi[\"<sos>\"]]\n        numericalized.extend(vocab.numericalize(sentence))\n        numericalized.append(vocab.stoi[\"<eos>\"])\n        return numericalized\n\n    def _load_from_file(self, img_path):\n        image = cv2.imread(img_path, cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image = cv2.resize(image, (IMG_SIZE, IMG_SIZE)) \n        image /= 255.0  # Normalize\n        return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def bms_collate(batch):\n    \n    imgs, label_formula, label_atom_conn, label_hydrogen, label_formula_len, label_atom_conn_len, label_hydrogen_len = [], [], [], [], [], [], []\n    \n    for data_point in batch:\n        imgs.append(data_point[0])\n        label_formula.append(data_point[1])\n        label_atom_conn.append(data_point[2])\n        label_hydrogen.append(data_point[3])\n        label_formula_len.append(data_point[4])\n        label_atom_conn_len.append(data_point[5])\n        label_hydrogen_len.append(data_point[6])\n\n    label_formula = pad_sequence(label_formula, batch_first=True, padding_value=vocab_formula.stoi[\"<pad>\"])\n    label_atom_conn = pad_sequence(label_atom_conn, batch_first=True, padding_value=vocab_atom_conn.stoi[\"<pad>\"])\n    label_hydrogen = pad_sequence(label_hydrogen, batch_first=True, padding_value=vocab_hydrogen.stoi[\"<pad>\"])\n\n    return torch.stack(imgs), label_formula, label_atom_conn, label_hydrogen, torch.stack(label_formula_len).reshape(-1, 1), torch.stack(label_atom_conn_len).reshape(-1, 1), torch.stack(label_hydrogen_len).reshape(-1, 1)\n    \n\ntrain_dataset = BMSDataset(train_df, get_train_augs(), vocab_formula, vocab_atom_conn, vocab_hydrogen)\ntrain_loader = DataLoader(\n    train_dataset,\n    batch_size=64,\n    pin_memory=True,\n    num_workers=4,\n    shuffle=True,\n    collate_fn=bms_collate\n)\n\nval_dataset = BMSDataset(test_df, get_valid_augs(), vocab_formula, vocab_atom_conn, vocab_hydrogen)\nval_loader = DataLoader(\n    val_dataset,\n    batch_size=64,\n    pin_memory=True,\n    num_workers=4,\n    shuffle=False,\n    collate_fn=bms_collate\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Available device: {device}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model Architecture"},{"metadata":{},"cell_type":"markdown","source":"This code is heavily based on this great tutorial: https://github.com/sgrvinod/a-PyTorch-Tutorial-to-Image-Captioning"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Encoder(nn.Module):\n    \"\"\"\n    Encoder.\n    \"\"\"\n\n    def __init__(self, back_bone, encoded_image_size=8):\n        super(Encoder, self).__init__()\n        self.enc_image_size = encoded_image_size\n\n        # Remove linear and pool layers (since we're not doing classification)\n        modules = list(back_bone.children())[:-2]\n        self.back_bone = nn.Sequential(*modules)\n\n        # Resize image to fixed size to allow input images of variable size\n        self.adaptive_pool = nn.AdaptiveAvgPool2d((encoded_image_size, encoded_image_size))\n\n        self.fine_tune(False)\n\n    def forward(self, images):\n        \"\"\"\n        Forward propagation.\n        :param images: images, a tensor of dimensions (batch_size, 3, image_size, image_size)\n        :return: encoded images\n        \"\"\"\n        # Extract features\n        out = self.back_bone(images)  # (batch_size, 2048, image_size/32, image_size/32)\n        out = self.adaptive_pool(out)  # (batch_size, 2048, encoded_image_size, encoded_image_size)\n        out = out.permute(0, 2, 3, 1)  # (batch_size, encoded_image_size, encoded_image_size, 2048)\n        return out\n\n    def fine_tune(self, fine_tune=True):\n        \"\"\"\n        Allow or prevent the computation of gradients for convolutional blocks 2 through 4 of the encoder.\n        :param fine_tune: Allow?\n        \"\"\"\n        for p in self.back_bone.parameters():\n            p.requires_grad = False\n        # If fine-tuning, only fine-tune convolutional blocks 2 through 4\n        for c in list(self.back_bone.children())[5:]:\n            for p in c.parameters():\n                p.requires_grad = fine_tune","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Attention(nn.Module):\n    \"\"\"\n    Attention Network.\n    \"\"\"\n\n    def __init__(self, encoder_dim, decoder_dim, attention_dim):\n        \"\"\"\n        :param encoder_dim: feature size of encoded images\n        :param decoder_dim: size of decoder's RNN\n        :param attention_dim: size of the attention network\n        \"\"\"\n        super(Attention, self).__init__()\n        self.encoder_att = nn.Linear(encoder_dim, attention_dim)  # linear layer to transform encoded image\n        self.decoder_att = nn.Linear(decoder_dim, attention_dim)  # linear layer to transform decoder's output\n        self.full_att = nn.Linear(attention_dim, 1)  # linear layer to calculate values to be softmax-ed\n        self.relu = nn.ReLU()\n        self.softmax = nn.Softmax(dim=1)  # softmax layer to calculate weights\n\n    def forward(self, encoder_out, decoder_hidden):\n        \"\"\"\n        Forward propagation.\n        :param encoder_out: encoded images, a tensor of dimension (batch_size, num_pixels, encoder_dim)\n        :param decoder_hidden: previous decoder output, a tensor of dimension (batch_size, decoder_dim)\n        :return: attention weighted encoding, weights\n        \"\"\"\n        att1 = self.encoder_att(encoder_out)  # (batch_size, num_pixels, attention_dim)\n        att2 = self.decoder_att(decoder_hidden)  # (batch_size, attention_dim)\n        att = self.full_att(self.relu(att1 + att2.unsqueeze(1))).squeeze(2)  # (batch_size, num_pixels)\n        alpha = self.softmax(att)  # (batch_size, num_pixels)\n        attention_weighted_encoding = (encoder_out * alpha.unsqueeze(2)).sum(dim=1)  # (batch_size, encoder_dim)\n\n        return attention_weighted_encoding, alpha","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DecoderWithAttention(nn.Module):\n    \"\"\"\n    Decoder.\n    \"\"\"\n\n    def __init__(self, attention_dim, embed_dim, decoder_dim, vocab_size, encoder_dim=2048, dropout=0.4):\n        \"\"\"\n        :param attention_dim: size of attention network\n        :param embed_dim: embedding size\n        :param decoder_dim: size of decoder's RNN\n        :param vocab_size: size of vocabulary\n        :param encoder_dim: feature size of encoded images\n        :param dropout: dropout\n        \"\"\"\n        super(DecoderWithAttention, self).__init__()\n\n        self.encoder_dim = encoder_dim\n        self.attention_dim = attention_dim\n        self.embed_dim = embed_dim\n        self.decoder_dim = decoder_dim\n        self.vocab_size = vocab_size\n        self.dropout = dropout\n\n        self.attention = Attention(encoder_dim, decoder_dim, attention_dim)  # attention network\n\n        self.embedding = nn.Embedding(vocab_size, embed_dim)  # embedding layer\n        self.dropout = nn.Dropout(p=self.dropout)\n        self.decode_step = nn.LSTMCell(embed_dim + encoder_dim, decoder_dim, bias=True)  # decoding LSTMCell\n        self.init_h = nn.Linear(encoder_dim, decoder_dim)  # linear layer to find initial hidden state of LSTMCell\n        self.init_c = nn.Linear(encoder_dim, decoder_dim)  # linear layer to find initial cell state of LSTMCell\n        self.f_beta = nn.Linear(decoder_dim, encoder_dim)  # linear layer to create a sigmoid-activated gate\n        self.sigmoid = nn.Sigmoid()\n        self.fc = nn.Linear(decoder_dim, vocab_size)  # linear layer to find scores over vocabulary\n        self.init_weights()  # initialize some layers with the uniform distribution\n\n    def init_weights(self):\n        \"\"\"\n        Initializes some parameters with values from the uniform distribution, for easier convergence.\n        \"\"\"\n        self.embedding.weight.data.uniform_(-0.1, 0.1)\n        self.fc.bias.data.fill_(0)\n        self.fc.weight.data.uniform_(-0.1, 0.1)\n\n    def load_pretrained_embeddings(self, embeddings):\n        \"\"\"\n        Loads embedding layer with pre-trained embeddings.\n        :param embeddings: pre-trained embeddings\n        \"\"\"\n        self.embedding.weight = nn.Parameter(embeddings)\n\n    def fine_tune_embeddings(self, fine_tune=True):\n        \"\"\"\n        Allow fine-tuning of embedding layer? (Only makes sense to not-allow if using pre-trained embeddings).\n        :param fine_tune: Allow?\n        \"\"\"\n        for p in self.embedding.parameters():\n            p.requires_grad = fine_tune\n\n    def init_hidden_state(self, encoder_out):\n        \"\"\"\n        Creates the initial hidden and cell states for the decoder's LSTM based on the encoded images.\n        :param encoder_out: encoded images, a tensor of dimension (batch_size, num_pixels, encoder_dim)\n        :return: hidden state, cell state\n        \"\"\"\n        mean_encoder_out = encoder_out.mean(dim=1)\n        h = self.init_h(mean_encoder_out)  # (batch_size, decoder_dim)\n        c = self.init_c(mean_encoder_out)\n        return h, c\n\n    def forward(self, encoder_out, encoded_captions, caption_lengths):\n        \"\"\"\n        Forward propagation.\n        :param encoder_out: encoded images, a tensor of dimension (batch_size, enc_image_size, enc_image_size, encoder_dim)\n        :param encoded_captions: encoded captions, a tensor of dimension (batch_size, max_caption_length)\n        :param caption_lengths: caption lengths, a tensor of dimension (batch_size, 1)\n        :return: scores for vocabulary, sorted encoded captions, decode lengths, weights, sort indices\n        \"\"\"\n        batch_size = encoder_out.size(0)\n        encoder_dim = encoder_out.size(-1)\n        vocab_size = self.vocab_size\n\n        # Flatten image\n        encoder_out = encoder_out.view(batch_size, -1, encoder_dim)  # (batch_size, num_pixels, encoder_dim)\n        num_pixels = encoder_out.size(1)\n\n        # Sort input data by decreasing lengths; why? apparent below\n        caption_lengths, sort_ind = caption_lengths.squeeze(1).sort(dim=0, descending=True)\n        encoder_out = encoder_out[sort_ind]\n        encoded_captions = encoded_captions[sort_ind]\n\n        # Embedding\n        embeddings = self.embedding(encoded_captions)  # (batch_size, max_caption_length, embed_dim)\n\n        # Initialize LSTM state\n        h, c = self.init_hidden_state(encoder_out)  # (batch_size, decoder_dim)\n\n        # We won't decode at the <end> position, since we've finished generating as soon as we generate <end>\n        # So, decoding lengths are actual lengths - 1\n        decode_lengths = (caption_lengths - 1).tolist()\n\n        # Create tensors to hold word predicion scores and alphas\n        predictions = torch.zeros(batch_size, max(decode_lengths), vocab_size).to(device)\n        alphas = torch.zeros(batch_size, max(decode_lengths), num_pixels).to(device)\n\n        # At each time-step, decode by\n        # attention-weighing the encoder's output based on the decoder's previous hidden state output\n        # then generate a new word in the decoder with the previous word and the attention weighted encoding\n        for t in range(max(decode_lengths)):\n            batch_size_t = sum([l > t for l in decode_lengths])\n            attention_weighted_encoding, alpha = self.attention(encoder_out[:batch_size_t],\n                                                                h[:batch_size_t])\n            gate = self.sigmoid(self.f_beta(h[:batch_size_t]))  # gating scalar, (batch_size_t, encoder_dim)\n            attention_weighted_encoding = gate * attention_weighted_encoding\n            h, c = self.decode_step(\n                torch.cat([embeddings[:batch_size_t, t, :], attention_weighted_encoding], dim=1),\n                (h[:batch_size_t], c[:batch_size_t]))  # (batch_size_t, decoder_dim)\n            preds = self.fc(self.dropout(h))  # (batch_size_t, vocab_size)\n            predictions[:batch_size_t, t, :] = preds\n            alphas[:batch_size_t, t, :] = alpha\n\n        return predictions, encoded_captions, decode_lengths, alphas, sort_ind","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"class AverageMeter:\n    \"\"\"\n    Keeps track of most recent, average, sum, and count of a metric.\n    \"\"\"\n\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def accuracy(scores, targets, k):\n    \"\"\"\n    Computes top-k accuracy, from predicted and true labels.\n    :param scores: scores from the model\n    :param targets: true labels\n    :param k: k in top-k accuracy\n    :return: top-k accuracy\n    \"\"\"\n\n    batch_size = targets.size(0)\n    _, ind = scores.topk(k, 1, True, True)\n    correct = ind.eq(targets.view(-1, 1).expand_as(ind))\n    correct_total = correct.view(-1).float().sum()  # 0D tensor\n    return correct_total.item() * (100.0 / batch_size)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def clip_gradient(optimizer, grad_clip):\n    \"\"\"\n    Clips gradients computed during backpropagation to avoid explosion of gradients.\n    :param optimizer: optimizer with the gradients to be clipped\n    :param grad_clip: clip value\n    \"\"\"\n    for group in optimizer.param_groups:\n        for param in group['params']:\n            if param.grad is not None:\n                param.grad.data.clamp_(-grad_clip, grad_clip)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_epoch(train_loader, encoder, decoders, criterion, encoder_optimizer, decoder_optimizers, epoch):\n    \"\"\"\n    Performs one epoch's training.\n    :param train_loader: DataLoader for training data\n    :param encoder: encoder model\n    :param decoder: decoder model\n    :param criterion: loss layer\n    :param encoder_optimizer: optimizer to update encoder's weights (if fine-tuning)\n    :param decoder_optimizer: optimizer to update decoder's weights\n    :param epoch: epoch number\n    \"\"\"\n    decoder_formula_optimizer, decoder_atom_conn_optimizer, decoder_hydrogen_optimizer = decoder_optimizers\n    decoder_formula, decoder_atom_conn, decoder_hydrogen = decoders\n    \n    batch_time = AverageMeter()  # forward prop. + back prop. time\n    losses_formula = AverageMeter()  # loss (per word decoded)\n    losses_atom_conn = AverageMeter()\n    losses_hydrogen = AverageMeter()\n    \n    # train mode (dropout and batchnorm is used)\n    decoder_formula.train(), decoder_atom_conn.train(), decoder_hydrogen.train()\n    start = time.time()\n    \n    for i, (imgs, label_formula, label_atom_conn, label_hydrogen, label_formula_len, label_atom_conn_len, label_hydrogen_lens) in tqdm(enumerate(train_loader), total=len(train_loader), position=0, leave=True):\n        # Move to GPU, if available\n        imgs = imgs.to(device)\n        label_formula = label_formula.to(device)\n        label_atom_conn = label_atom_conn.to(device)\n        label_hydrogen = label_hydrogen.to(device)\n        \n        label_formula_len = label_formula_len.to(device)\n        label_atom_conn_len = label_atom_conn_len.to(device)\n        label_hydrogen_lens = label_hydrogen_lens.to(device)\n\n        # Forward prop.\n        imgs = encoder(imgs)\n        scores_formula, caps_sorted_formula, decode_lengths_formula, alphas_formula, sort_ind_formula = decoder_formula(imgs, label_formula, label_formula_len)\n        scores_atom_conn, caps_sorted_atom_conn, decode_lengths_atom_conn, alphas_atom_conn, sort_ind_atom_conn = decoder_atom_conn(imgs, label_atom_conn, label_atom_conn_len)\n        scores_hydrogen, caps_sorted_hydrogen, decode_lengths_hydrogen, alphas_hydrogen, sort_ind_hydrogen = decoder_hydrogen(imgs, label_hydrogen, label_hydrogen_lens)\n\n        # Since we decoded starting with <start>, the targets are all words after <start>, up to <end>\n        targets_formula = caps_sorted_formula[:, 1:]\n        targets_atom_conn = caps_sorted_atom_conn[:, 1:]\n        targets_hydrogen = caps_sorted_hydrogen[:, 1:]\n\n        # Remove timesteps that we didn't decode at, or are pads\n        # pack_padded_sequence is an easy trick to do this\n        scores_formula = pack_padded_sequence(scores_formula, decode_lengths_formula, batch_first=True).data\n        targets_formula = pack_padded_sequence(targets_formula, decode_lengths_formula, batch_first=True).data\n\n        \n        scores_atom_conn = pack_padded_sequence(scores_atom_conn, decode_lengths_atom_conn, batch_first=True).data\n        targets_atom_conn = pack_padded_sequence(targets_atom_conn, decode_lengths_atom_conn, batch_first=True).data\n        \n        scores_hydrogen = pack_padded_sequence(scores_hydrogen, decode_lengths_hydrogen, batch_first=True).data\n        targets_hydrogen = pack_padded_sequence(targets_hydrogen, decode_lengths_hydrogen, batch_first=True).data\n        \n        # Calculate loss\n        loss_formula = criterion(scores_formula, targets_formula)\n        loss_atom_conn = criterion(scores_atom_conn, targets_atom_conn)\n        loss_hydrogen = criterion(scores_hydrogen, targets_hydrogen)\n\n        # Add doubly stochastic attention regularization\n        loss_formula += alpha_c * ((1. - alphas_formula.sum(dim=1)) ** 2).mean()\n        loss_atom_conn += alpha_c * ((1. - alphas_atom_conn.sum(dim=1)) ** 2).mean()\n        loss_hydrogen += alpha_c * ((1. - alphas_hydrogen.sum(dim=1)) ** 2).mean()\n\n        # Back prop.\n        decoder_formula_optimizer.zero_grad()\n        decoder_atom_conn_optimizer.zero_grad()\n        decoder_hydrogen_optimizer.zero_grad()\n        \n        if encoder_optimizer is not None:\n            encoder_optimizer.zero_grad()\n        \n        # Backpropagation\n        loss_formula.backward()\n        loss_atom_conn.backward()\n        loss_hydrogen.backward()\n\n        # Clip gradients\n        if grad_clip is not None:\n            clip_gradient(decoder_formula_optimizer, grad_clip)\n            clip_gradient(decoder_atom_conn_optimizer, grad_clip)\n            clip_gradient(decoder_hydrogen_optimizer, grad_clip)\n\n            if encoder_optimizer is not None:\n                clip_gradient(encoder_optimizer, grad_clip)\n\n        # Update weights\n        decoder_formula_optimizer.step()\n        decoder_atom_conn_optimizer.step()\n        decoder_hydrogen_optimizer.step()\n        if encoder_optimizer is not None:\n            encoder_optimizer.step()\n\n        # Keep track of metrics\n        losses_formula.update(loss_formula.item(), sum(decode_lengths_formula))\n        losses_atom_conn.update(loss_atom_conn.item(), sum(decode_lengths_atom_conn))\n        losses_hydrogen.update(loss_hydrogen.item(), sum(decode_lengths_hydrogen))\n        batch_time.update(time.time() - start)\n\n        start = time.time()\n        \n    return losses_formula, losses_atom_conn, losses_hydrogen, batch_time","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def validate_epoch(val_loader, encoder, decoders, criterion):\n    \n    decoder_formula, decoder_atom_conn, decoder_hydrogen = decoders\n    \n    decoder_formula.eval()  # eval mode (no dropout or batchnorm)\n    decoder_atom_conn.eval()\n    decoder_hydrogen.eval()\n\n    if encoder is not None:\n        encoder.eval()\n    \n    batch_time = AverageMeter()\n    losses_formula = AverageMeter()  # loss (per word decoded)\n    losses_atom_conn = AverageMeter()\n    losses_hydrogen = AverageMeter()\n    \n    start = time.time()\n\n    references = list()  # references (true captions) for calculating BLEU-4 score\n    hypotheses = list()  # hypotheses (predictions)\n\n    with torch.no_grad():\n        # Batches\n        for i, (imgs, label_formula, label_atom_conn, label_hydrogen, label_formula_len, label_atom_conn_len, label_hydrogen_lens) in tqdm(enumerate(val_loader), total=len(val_loader), position=0, leave=True):\n\n            # Move to device, if available\n            imgs = imgs.to(device)\n            label_formula = label_formula.to(device)\n            label_atom_conn = label_atom_conn.to(device)\n            label_hydrogen = label_hydrogen.to(device)\n\n            label_formula_len = label_formula_len.to(device)\n            label_atom_conn_len = label_atom_conn_len.to(device)\n            label_hydrogen_lens = label_hydrogen_lens.to(device)\n\n            # Forward prop.\n            if encoder is not None:\n                imgs = encoder(imgs)\n            \n            scores_formula, caps_sorted_formula, decode_lengths_formula, alphas_formula, sort_ind_formula = decoder_formula(imgs, label_formula, label_formula_len)\n            scores_atom_conn, caps_sorted_atom_conn, decode_lengths_atom_conn, alphas_atom_conn, sort_ind_atom_conn = decoder_atom_conn(imgs, label_atom_conn, label_atom_conn_len)\n            scores_hydrogen, caps_sorted_hydrogen, decode_lengths_hydrogen, alphas_hydrogen, sort_ind_hydrogen = decoder_hydrogen(imgs, label_hydrogen, label_hydrogen_lens)\n\n            # Since we decoded starting with <start>, the targets are all words after <start>, up to <end>\n            targets_formula = caps_sorted_formula[:, 1:]\n            targets_atom_conn = caps_sorted_atom_conn[:, 1:]\n            targets_hydrogen = caps_sorted_hydrogen[:, 1:]\n\n            # Remove timesteps that we didn't decode at, or are pads\n            # pack_padded_sequence is an easy trick to do this\n            scores_formula = pack_padded_sequence(scores_formula, decode_lengths_formula, batch_first=True).data\n            targets_formula = pack_padded_sequence(targets_formula, decode_lengths_formula, batch_first=True).data\n\n            scores_atom_conn = pack_padded_sequence(scores_atom_conn, decode_lengths_atom_conn, batch_first=True).data\n            targets_atom_conn = pack_padded_sequence(targets_atom_conn, decode_lengths_atom_conn, batch_first=True).data\n\n            scores_hydrogen = pack_padded_sequence(scores_hydrogen, decode_lengths_hydrogen, batch_first=True).data\n            targets_hydrogen = pack_padded_sequence(targets_hydrogen, decode_lengths_hydrogen, batch_first=True).data\n\n            # Calculate loss\n            loss_formula = criterion(scores_formula, targets_formula)\n            loss_atom_conn = criterion(scores_atom_conn, targets_atom_conn)\n            loss_hydrogen = criterion(scores_hydrogen, targets_hydrogen)\n\n            # Add doubly stochastic attention regularization\n            loss_formula += alpha_c * ((1. - alphas_formula.sum(dim=1)) ** 2).mean()\n            loss_atom_conn += alpha_c * ((1. - alphas_atom_conn.sum(dim=1)) ** 2).mean()\n            loss_hydrogen += alpha_c * ((1. - alphas_hydrogen.sum(dim=1)) ** 2).mean()\n\n\n            # Keep track of metrics\n            losses_formula.update(loss_formula.item(), sum(decode_lengths_formula))\n            losses_atom_conn.update(loss_atom_conn.item(), sum(decode_lengths_atom_conn))\n            losses_hydrogen.update(loss_hydrogen.item(), sum(decode_lengths_hydrogen))\n            batch_time.update(time.time() - start)\n            start = time.time()\n        \n        return losses_formula, losses_atom_conn, losses_hydrogen, batch_time","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Model parameters\nemb_dim = 512  # dimension of word embeddings\nattention_dim = 512  # dimension of attention linear layers\ndecoder_dim = 512  # dimension of decoder RNN\ndropout = 0.4\ncudnn.benchmark = False  # set to true only if inputs to model are fixed size; otherwise lot of computational overhead\n\n# Training parameters\nstart_epoch = 0\nepochs = 20  # number of epochs to train for (if early stopping is not triggered)\nepochs_since_improvement = 0  # keeps track of number of epochs since there's been an improvement in validation BLEU\nworkers = 4  # for data-loading; right now, only 1 works with h5py\nencoder_lr = 2e-4  # learning rate for encoder if fine-tuning\ndecoder_lr = 2e-4  # learning rate for decoder\ngrad_clip = 5.  # clip gradients at an absolute value of\nalpha_c = 1.  # regularization parameter for 'doubly stochastic attention', as in the paper\nbest_bleu4 = 0.  # BLEU-4 score right now\nfine_tune_encoder = False  # fine-tune encoder?","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Train 3 separate decoders to reduce the sequence length while predicting\ndecoder_formula = DecoderWithAttention(\n    attention_dim=attention_dim,\n    embed_dim=emb_dim,\n    decoder_dim=decoder_dim,\n    vocab_size=len(vocab_formula),\n    dropout=dropout\n)\ndecoder_atom_conn = DecoderWithAttention(\n    attention_dim=attention_dim,\n    embed_dim=emb_dim,\n    decoder_dim=decoder_dim,\n    vocab_size=len(vocab_atom_conn),\n    dropout=dropout\n)\ndecoder_hydrogen = DecoderWithAttention(\n    attention_dim=attention_dim,\n    embed_dim=emb_dim,\n    decoder_dim=decoder_dim,\n    vocab_size=len(vocab_hydrogen),\n    dropout=dropout\n)\n\ndecoder_formula_optimizer = torch.optim.Adam(\n    params=filter(lambda p: p.requires_grad, decoder_formula.parameters()),\n    lr=decoder_lr\n)\n\ndecoder_atom_conn_optimizer = torch.optim.Adam(\n    params=filter(lambda p: p.requires_grad, decoder_atom_conn.parameters()),\n    lr=decoder_lr\n)\n\ndecoder_hydrogen_optimizer = torch.optim.Adam(\n    params=filter(lambda p: p.requires_grad, decoder_hydrogen.parameters()),\n    lr=decoder_lr\n)\n\nencoder = Encoder(timm.create_model('resnet101', pretrained=True))\nencoder.fine_tune(fine_tune_encoder)\nencoder_optimizer = torch.optim.Adam(\n    params=filter(lambda p: p.requires_grad, encoder.parameters()),\n    lr=encoder_lr\n) if fine_tune_encoder else None","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"encoder = encoder.to(device)\ndecoder_formula = decoder_formula.to(device)\ndecoder_atom_conn = decoder_atom_conn.to(device)\ndecoder_hydrogen = decoder_hydrogen.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Loss function\ncriterion = nn.CrossEntropyLoss().to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def save_checkpoint(epoch, encoder, decoder_formula, decoder_atom_conn, decoder_hydrogen, encoder_optimizer, decoder_formula_optimizer, decoder_atom_conn_optimizer, decoder_hydrogen_optimizer):\n    \"\"\"\n    Saves model checkpoint.\n    :param data_name: base name of processed dataset\n    :param epoch: epoch number\n    :param epochs_since_improvement: number of epochs since last improvement in BLEU-4 score\n    :param encoder: encoder model\n    :param decoder: decoder model\n    :param encoder_optimizer: optimizer to update encoder's weights, if fine-tuning\n    :param decoder_optimizer: optimizer to update decoder's weights\n    :param bleu4: validation BLEU-4 score for this epoch\n    :param is_best: is this checkpoint the best so far?\n    \"\"\"\n    state = {\n        'epoch': epoch,\n        'encoder': encoder,\n        'decoder_formula': decoder_formula,\n        'decoder_atom_conn': decoder_atom_conn,\n        'decoder_hydrogen': decoder_hydrogen,\n        'encoder_optimizer': encoder_optimizer,\n        'decoder_formula_optimizer': decoder_formula_optimizer,\n        'decoder_atom_conn_optimizer': decoder_atom_conn_optimizer,\n        'decoder_hydrogen_optimizer': decoder_hydrogen_optimizer\n    }\n    filename = f'checkpoint_{epoch}.pth.tar'\n    torch.save(state, filename)\n    # If this checkpoint is the best so far, store a copy so it doesn't get overwritten by a worse checkpoint","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# For inference:\nepochs = 0","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"best_score = float('inf')\nfor epoch in range(start_epoch, epochs):\n\n#     # Decay learning rate if there is no improvement for 8 consecutive epochs, and terminate training after 20\n#     if epochs_since_improvement == 20:\n#         break\n#     if epochs_since_improvement > 0 and epochs_since_improvement % 8 == 0:\n#         adjust_learning_rate(decoder_optimizer, 0.8)\n#         if fine_tune_encoder:\n#             adjust_learning_rate(encoder_optimizer, 0.8)\n\n    # One epoch's training\n    train_losses_formula, train_losses_atom_conn, train_losses_hydrogen, batch_time = train_epoch(\n        train_loader=train_loader,\n        encoder=encoder,\n        decoders=(decoder_formula, decoder_atom_conn, decoder_hydrogen),\n        criterion=criterion,\n        encoder_optimizer=encoder_optimizer,\n        decoder_optimizers=(decoder_formula_optimizer, decoder_atom_conn_optimizer, decoder_hydrogen_optimizer),\n        epoch=epoch\n    )\n    \n    val_losses_formula, val_losses_atom_conn, val_losses_hydrogen, _ = validate_epoch(\n        val_loader,\n        encoder,\n        decoders=(decoder_formula, decoder_atom_conn, decoder_hydrogen),\n        criterion=criterion\n    )\n    \n    if best_score > (val_losses_formula.avg + val_losses_atom_conn.avg + val_losses_hydrogen.avg)/3:\n        best_score = (val_losses_formula.avg + val_losses_atom_conn.avg + val_losses_hydrogen.avg)/3\n        print(f\"Saving checkpoint. Best score: {best_score:.4f}\")\n        save_checkpoint(\n            epoch+1,\n            encoder,\n            decoder_formula,\n            decoder_atom_conn,\n            decoder_hydrogen,\n            encoder_optimizer,\n            decoder_formula_optimizer,\n            decoder_atom_conn_optimizer,\n            decoder_hydrogen_optimizer\n        )\n    \n    print(f'Epoch: {epoch+1:02} | Time: {batch_time.avg} sec\\n')\n    print(f'\\t    Train Losses (Avg.): {(train_losses_formula.avg + train_losses_atom_conn.avg + train_losses_hydrogen.avg)/3:.4f}')\n    print(f'\\t    \\tFormula: {train_losses_formula.avg:.4f} | Atom Connections: {train_losses_atom_conn.avg:.4f} | Hydrogen: {train_losses_hydrogen.avg:.4f}')\n    print(f'\\t    Validation Losses (Avg.): {(val_losses_formula.avg + val_losses_atom_conn.avg + val_losses_hydrogen.avg)/3:.4f}')\n    print(f'\\t    \\tFormula: {val_losses_formula.avg:.4f} | Atom Connections: {val_losses_atom_conn.avg:.4f} | Hydrogen: {val_losses_hydrogen.avg:.4f}\\n')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Inference"},{"metadata":{"trusted":true},"cell_type":"code","source":"checkpoint = torch.load('../input/bms-mt-show-attend-and-tell-pytorch-baseline/checkpoint_20.pth.tar')\nencoder = checkpoint['encoder']\ndecoder_formula = checkpoint['decoder_formula']\ndecoder_atom_conn = checkpoint['decoder_atom_conn']\ndecoder_hydrogen = checkpoint['decoder_hydrogen']","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_sub_df = pd.read_csv(\"../input/bms-molecular-translation/sample_submission.csv\")\nsample_sub_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class BMSDatasetTest(Dataset):\n    def __init__(self, df):\n        super().__init__()\n        self.df = df\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        # Read the image\n        img_id = self.df.iloc[idx]['image_id']\n        img_path = os.path.join('../input/bms-molecular-translation/test', img_id[0], img_id[1], img_id[2], f'{img_id}.png')\n        \n        img = self._load_from_file(img_path)\n        return torch.tensor(img).permute(2, 0, 1)\n\n    def _load_from_file(self, img_path):\n        image = cv2.imread(img_path, cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image = cv2.resize(image, (IMG_SIZE, IMG_SIZE)) \n        image /= 255.0  # Normalize\n        return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_dataset = BMSDatasetTest(sample_sub_df)\ntest_batch_size = 128\ntest_dataloader = DataLoader(\n    test_dataset,\n    batch_size=test_batch_size,\n    pin_memory=True,\n    num_workers=4,\n    shuffle=False,\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"max_pred_len_formula = int(np.mean(label_df['formula'].str.len()))+1\nmax_pred_len_atom_conn = int(np.mean(label_df['atom_conn'].str.len()))+1\nmax_pred_len_hydrogen = int(np.mean(label_df['hydrogen'].str.len()))+1\n\nprint(f'Max formula pred. length: {max_pred_len_formula}')\nprint(f'Max atom conn. pred. length: {max_pred_len_atom_conn}')\nprint(f'Max hydrogen pred. length: {max_pred_len_hydrogen}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def decode(decoder, vocab, encoder_out, batch_size, max_pred_len):\n    h, c = decoder.init_hidden_state(encoder_out)\n    start = torch.full((batch_size, 1), vocab.stoi['<sos>']).to(device)\n    pred = torch.zeros((batch_size, max_pred_len), dtype=torch.long).to(device)\n    pred[:, 0] = start.squeeze()\n    \n    idx = 1\n\n    while True:\n        embeddings = decoder.embedding(start).squeeze(1)\n        \n        awe, _ = decoder.attention(encoder_out, h)\n        \n        gate = decoder.sigmoid(decoder.f_beta(h))  # gating scalar, (s, encoder_dim)\n        awe = gate * awe\n\n        h, c = decoder.decode_step(torch.cat([embeddings, awe], dim=1), (h, c))\n        scores = decoder.fc(h)  # (s, vocab_size)\n        scores = F.log_softmax(scores, dim=1)\n\n        start = scores.argmax(1).reshape(-1, 1).to(device)\n        \n        pred[:, idx] = start.squeeze(1)\n        \n        if idx >= max_pred_len-1:\n            break\n        \n        idx += 1\n        \n    return pred","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def inference(encoder, decoders, imgs, vocabs):\n    imgs = imgs.to(device)\n    batch_size = len(imgs)\n    \n    decoder_formula, decoder_atom_conn, decoder_hydrogen = decoders\n    vocab_formula, vocab_atom_conn, vocab_hydrogen = vocabs\n    \n    encoder_out = encoder(imgs)  # (1, enc_image_size, enc_image_size, encoder_dim)\n    enc_image_size = encoder_out.size(1)\n    encoder_dim = encoder_out.size(3)\n    \n    encoder_out = encoder_out.view(batch_size, -1, encoder_dim)\n    \n    ####################\n    # Formula Decoding\n    ####################\n    pred_formula = decode(decoder_formula, vocab_formula, encoder_out, batch_size, max_pred_len_formula)\n    pred_atom_conn = decode(decoder_atom_conn, vocab_atom_conn, encoder_out, batch_size, max_pred_len_atom_conn)\n    pred_hydrogen = decode(decoder_hydrogen, vocab_hydrogen, encoder_out, batch_size, max_pred_len_hydrogen)\n    \n    return pred_formula, pred_atom_conn, pred_hydrogen","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def batch_stringify(batch, vocab):\n    preds = []\n    for item in batch:\n        pred = np.vectorize(vocab.itos.get)(item)\n        # Truncate everything after <eos>\n        try:\n            pred = pred[1:np.nonzero(pred == '<eos>')[0][0]]\n        except IndexError:\n            pred = pred[1: ]\n            pass\n\n        preds.append(\"\".join(pred))\n    return preds","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Calculating Competition Metric on validation data"},{"metadata":{"trusted":true},"cell_type":"code","source":"decoders = (decoder_formula, decoder_atom_conn, decoder_hydrogen)\nvocabs = (vocab_formula, vocab_atom_conn, vocab_hydrogen)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"preds_formula, preds_atom_conn, preds_hydrogen = [], [], []\ngts_formula, gts_atom_conn, gts_hydrogen = [], [], []\n\nfor imgs, label_formula, label_atom_conn, label_hydrogen, label_formula_len, label_atom_conn_len, label_hydrogen_lens in tqdm(val_loader):\n    pred_formula, pred_atom_conn, pred_hydrogen = inference(encoder, decoders, imgs, vocabs)\n    \n    pred_formula = pred_formula.cpu().detach().numpy()\n    pred_atom_conn = pred_atom_conn.cpu().detach().numpy()\n    pred_hydrogen = pred_hydrogen.cpu().detach().numpy()\n\n    preds_formula.extend(batch_stringify(pred_formula, vocab_formula))\n    preds_atom_conn.extend(batch_stringify(pred_atom_conn, vocab_atom_conn))\n    preds_hydrogen.extend(batch_stringify(pred_hydrogen, vocab_hydrogen))\n    \n    gts_formula.extend(batch_stringify(label_formula, vocab_formula))\n    gts_atom_conn.extend(batch_stringify(label_atom_conn, vocab_atom_conn))\n    gts_hydrogen.extend(batch_stringify(label_hydrogen, vocab_hydrogen))\n    \n\nlev_dist_formula = np.mean(np.vectorize(levenshtein_distance)(preds_formula, gts_formula))\nlev_dist_atom_conn = np.mean(np.vectorize(levenshtein_distance)(preds_atom_conn, gts_atom_conn))\nlev_dist_hydrogen = np.mean(np.vectorize(levenshtein_distance)(preds_hydrogen, gts_hydrogen))\n\nprint(f'Levenshtein distance (Formula): {lev_dist_formula}')\nprint(f'Levenshtein distance (Atom Conn.): {lev_dist_atom_conn}')\nprint(f'Levenshtein distance (Hydrogen): {lev_dist_hydrogen}\\n')\nprint(f'Avg. Validation Levenshtein distance {(lev_dist_formula + lev_dist_atom_conn + lev_dist_hydrogen)/3}')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prediction on Test Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"preds_formula, preds_atom_conn, preds_hydrogen = [], [], []\nfor imgs in tqdm(test_dataloader):\n    pred_formula, pred_atom_conn, pred_hydrogen = inference(encoder, decoders, imgs, vocabs)\n    \n    pred_formula = pred_formula.cpu().detach().numpy()\n    pred_atom_conn = pred_atom_conn.cpu().detach().numpy()\n    pred_hydrogen = pred_hydrogen.cpu().detach().numpy()\n    \n    preds_formula.extend(batch_stringify(pred_formula, vocab_formula))\n    preds_atom_conn.extend(batch_stringify(pred_atom_conn, vocab_atom_conn))\n    preds_hydrogen.extend(batch_stringify(pred_hydrogen, vocab_hydrogen))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_sub_df['formula'] = preds_formula\nsample_sub_df['atom_conn'] = preds_atom_conn\nsample_sub_df['hydrogen'] = preds_hydrogen","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_sub_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_sub_df['InChI'] = sample_sub_df[['formula', 'atom_conn', 'hydrogen']].apply(lambda row: f'InChI=1S/{row[\"formula\"]}/c{row[\"atom_conn\"]}/h{row[\"hydrogen\"]}', axis=1)\nsample_sub_df = sample_sub_df.drop(['formula', 'atom_conn', 'hydrogen'], axis=1)\nsample_sub_df.to_csv(\"submission.csv\", index=False)\nsample_sub_df.head()","execution_count":null,"outputs":[]}],"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}