{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Top"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport re\nimport sys, time, os, gc\nimport cv2\nimport numba\nimport torch\nimport torchvision\nfrom Levenshtein import distance\n\nfrom torch.nn.utils.rnn import pad_sequence\nfrom torch.utils.data import Dataset, DataLoader\nfrom albumentations.pytorch import ToTensorV2\nimport albumentations as A\nimport torch.nn as nn\n\nfrom sklearn.model_selection import train_test_split as tts\nfrom tqdm.notebook import tqdm\n\ntqdm().pandas()\n\nSEED = 1023\n\nnp.random.seed(SEED)\n\nstart_time = time.time()\ntime_limit = 8.5 * 3600 #8.5 hours limit","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Read Datasets"},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv(\"../input/bms-molecular-translation/train_labels.csv\")\nprint(train.shape)\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test = pd.read_csv(\"../input/bms-molecular-translation/sample_submission.csv\")\nprint(test.shape)\ntest.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Define Constants"},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN = True #To control whether to train the model or not\nPREDICT = False #To control whether to predict submission or not\nDEBUG = True\n\nIMG_SIZE = 224\nMAX_LEN = 400\nVOCAB_SIZE = 40\nBATCH_SIZE = 128\nWORKERS = 4\nEPOCHS = 3\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Current device:\", DEVICE.upper())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if DEBUG:\n    sample_factor = 0.20\n    n = int(train.shape[0]*sample_factor)\n    train = train.sample(n, random_state=SEED*3).reset_index(drop=True)\n    print(\"New train size:\", train.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"VOCAB = ['<pad>', '<sos>', '<eos>',\n         '(', ')', '+', ',', '-', '=', '/b', '/c', '/h', '/i', '/m', '/s', '/t',\n         *[str(x) for x in range(10)],\n         'B', 'Br', 'C', 'Cl', 'D', 'F', 'H', 'I', 'N', 'O', 'P', 'S', 'Si', 'T']\n\nstoi = {b:n for n,b in enumerate(VOCAB)}\nitos = {n:b for n,b in enumerate(VOCAB)}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Helper Functions"},{"metadata":{"trusted":true},"cell_type":"code","source":"def bms_collate(batch):\n    images, labels = [], []\n    for data in batch:\n        images.append(data[0])\n        labels.append(data[1])\n    labels = pad_sequence(labels, batch_first = True, padding_value = 0) #0 is the value of stoi[\"<pad>\"]\n    return torch.stack(images), labels\n\ndef generate_inchi(pred):\n    label = [itos[i] for i in pred.to(\"cpu\").numpy()]\n    result = []\n    for i in range(len(label)):\n        if label[i] == \"<eos>\":\n            break\n        result.append(label[i])\n    result = \"InChI=1S/\" + \"\".join(result)\n    return result\n\ndef levenshtein_score(y_pred, y):\n    #y_pred = np.argmax(y_pred, axis=2).astype(uint8)\n    score = []\n    for i in range(len(y)):\n        a = generate_inchi(y_pred[i])\n        b = generate_inchi(y[i])\n        #print(a, b, sep='\\n')\n        score.append(distance(a, b))\n    return score","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataloader"},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImageDataLoader(Dataset):\n    '''Loads the image dataset and prepares it for feedint to model.'''\n    \n    def __init__(self, df, img_size=224, train=True):\n        '''\n        df         => DataFrame where data will be extracted\n        paths      => List of paths to the image files\n        inchi      => List of InChI strings\n        inchi_pattern => RegEx pattern for tokenizing InChI\n        vocab      => List of all tokens to be used\n        stoi       => Converts string token to integer\n        itos       => Converts integer token to string\n        '''\n        self.df = df\n        self.img_size = img_size\n        self.is_train = train\n        loc = \"train\" if self.is_train else \"test\"\n        self.paths = df[\"image_id\"].apply(\n            lambda x: f\"../input/bms-molecular-translation/{loc}/{x[0]}/{x[1]}/{x[2]}/{x}.png\"\n        )\n        \n        self.inchi = self.df[\"InChI\"].values\n        self.inchi_pattern = r\"([A-Z][a-z]?|[0-9]|\\/[a-z]?|[\\S])\"\n        \n    def __len__(self):\n        return self.df.shape[0]\n    \n    def __getitem__(self, idx):\n        '''Fetch image at index idx. If train, fetch the labels.'''\n        \n        path = self.paths[idx]\n        image = cv2.imread(path)\n        #Read Image\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        #Flip: Majority of Image is white, flip white and black\n        image = np.apply_along_axis(lambda x: 255.0 - x, 0, image)\n        image = self.rotate_if_not_landscape(image)\n        image = A.Compose([\n            A.Resize(self.img_size, self.img_size),\n            A.Normalize(),\n            ToTensorV2()\n        ])(image = image)[\"image\"]\n        \n        if self.is_train:\n            inchi = self.inchi[idx]\n            inchi = [stoi[s] for s in self.tokenize_inchi(inchi)]\n            inchi.insert(0, stoi[\"<sos>\"])\n            inchi.append(stoi[\"<eos>\"])\n            inchi = torch.LongTensor(inchi)\n            return image, inchi\n        else:\n            return image\n    \n    def tokenize_inchi(self, inchi):\n        '''Aply Regex pattern to split InChI string.'''\n        #Since the part \"InChI=1S/\" is similar across all InChI, remove it\n        # then apply regex using the pattern to search for:\n        # Element, Number, Separator, and any other characters\n        return re.findall(self.inchi_pattern, inchi[9:])\n    \n    def rotate_if_not_landscape(self, image):\n        '''If the image is long vertically (portrait), rotate 90 deg.'''\n        if image.shape[0] < image.shape[1]:\n            return np.rot90(image, axes=(0, 1))\n        return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_train, df_valid = tts(train, test_size = 0.1, shuffle = True)\ndf_train, df_valid = df_train.reset_index(drop=True), df_valid.reset_index(drop=True)\nprint(\"Train:\", df_train.shape)\nprint(\"Validation:\", df_valid.shape)\n\ntrain_data = ImageDataLoader(df_train)\nvalid_data = ImageDataLoader(df_valid)\ntrain_loader = DataLoader(train_data, batch_size = BATCH_SIZE, shuffle = True,\n                          drop_last = True, collate_fn = bms_collate, num_workers = WORKERS,\n                          prefetch_factor=BATCH_SIZE//WORKERS)\n\nvalid_loader = DataLoader(valid_data, batch_size = BATCH_SIZE * 2, shuffle = False, drop_last = False,\n                          collate_fn = bms_collate, num_workers = WORKERS,\n                          prefetch_factor=(2*BATCH_SIZE)//WORKERS)\n\ndf_test = test.reset_index(drop=True)\nprint(\"Test:\", df_test.shape)\ntest_data = ImageDataLoader(df_test, train=False)\ntest_loader = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=False,\n                         drop_last = False, num_workers = WORKERS,\n                         prefetch_factor=BATCH_SIZE//WORKERS)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{},"cell_type":"markdown","source":"## Encoder"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Encoder(nn.Module):\n    '''Encodes the data input into a new vectorized representation.'''\n    def __init__(self):\n        super().__init__()\n        self.resnet = torchvision.models.resnet18(pretrained=True)\n        num_ftrs = self.resnet.fc.in_features\n        self.resnet.fc = nn.Linear(num_ftrs, 256)\n        self.resnet = self.resnet.to(DEVICE)\n        \n    def forward(self, images):\n        batch_size = images.size(0)\n        features = self.resnet(images)\n        features = features.permute(0, 1)\n        features = features.view(features.size(0), -1, features.size(-1))\n        return features","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Decoder"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Decoder(nn.Module):\n    def __init__(self, vocab_size, embed_size, encoder_dim, decoder_dim):\n        super().__init__()\n        self.vocab_size = vocab_size\n        self.decoder_dim = decoder_dim\n        self.embedding = nn.Embedding(vocab_size, embed_size)\n        self.init_h = nn.Linear(in_features = encoder_dim, out_features = decoder_dim)\n        self.init_c = nn.Linear(in_features = encoder_dim, out_features = decoder_dim)\n        self.lstm = nn.LSTMCell(input_size = (embed_size + encoder_dim), hidden_size = decoder_dim, bias = True)\n        self.drop = nn.Dropout(p = 0.3)\n        self.linear = nn.Linear(in_features = decoder_dim, out_features = vocab_size)\n        \n    def forward(self, features, inchis):\n        embeds = self.embedding(inchis)\n        \n        features = features.mean(dim = 1)\n        h = self.init_h(features)\n        c = self.init_c(features)\n        \n        seq_length = len(inchis[0]) - 1\n        batch_size = inchis.size(0)\n        preds = torch.zeros(batch_size, seq_length, self.vocab_size).to(DEVICE)\n        \n        for s in range(seq_length):\n            lstm_input = torch.cat((embeds[:, s], features), dim = 1)\n            h, c = self.lstm(lstm_input, (h, c))\n            x = self.drop(h)\n            x = self.linear(x)\n            preds[:, s] = x\n        return preds\n\n    def predict(self, features, max_len):\n        batch_size = features.size(0)\n        features = features.mean(dim = 1)\n        h = self.init_h(features)\n        c = self.init_c(features)\n        \n        word = torch.full((batch_size, 1), stoi[\"<sos>\"]).to(DEVICE)\n        embeds = self.embedding(word)\n        preds = torch.zeros((batch_size, max_len), dtype = torch.long).to(DEVICE)\n        preds[:, 0] = word.squeeze()\n        for i in range(max_len):\n            lstm_input = torch.cat((embeds[:, 0], features), dim = 1)\n            h, c = self.lstm(lstm_input, (h, c))\n            x = self.drop(h)\n            x = self.linear(x)\n            x = x.view(batch_size, -1)\n            pred_idx = x.argmax(dim = 1)\n            preds[:, i] = pred_idx\n            embeds = self.embedding(pred_idx).unsqueeze(1)\n        return preds","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train"},{"metadata":{"trusted":true},"cell_type":"code","source":"encoder = Encoder().to(DEVICE)\ndecoder = Decoder(\n    vocab_size = VOCAB_SIZE, embed_size = 256, encoder_dim = 256, decoder_dim = 512\n).to(DEVICE)\n\nencoder_optimizer = torch.optim.Adam(\n    encoder.parameters(), lr = 1e-4, weight_decay = 1e-6, amsgrad = False)\ndecoder_optimizer = torch.optim.Adam(\n    decoder.parameters(), lr = 1e-4, weight_decay = 1e-6, amsgrad = True)\n\nencoder_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n    encoder_optimizer, T_max = 4, eta_min = 1e-6, last_epoch = -1 )\ndecoder_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n    decoder_optimizer, T_max = 4, eta_min = 1e-6, last_epoch = -1 )\n\ncriterion = nn.CrossEntropyLoss(ignore_index = 0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train Loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"if TRAIN:\n    #Import trained weights\n    encoder.load_state_dict(torch.load(\n        \"../input/pytorch-train-with-lstm/bms_encoder.pth\", map_location = DEVICE))\n    decoder.load_state_dict(torch.load(\n        \"../input/pytorch-train-with-lstm/bms_decoder.pth\", map_location = DEVICE))\n    best_loss = np.inf\n    for epoch in range(EPOCHS):\n        print(f\"Epoch [{epoch+1}/{EPOCHS}]:\")\n        \n        if time.time() - start_time > time_limit:\n            break #Avoid TimeLimitExceeded Error\n            \n        encoder.train()\n        decoder.train()\n        for images, inchis in tqdm(train_loader):\n            images = images.to(DEVICE)\n            inchis = inchis.to(DEVICE)\n            encoder_optimizer.zero_grad()\n            decoder_optimizer.zero_grad()\n            encoded = encoder(images)\n            preds = decoder(encoded, inchis)\n            loss = criterion(preds.permute(0, 2, 1), inchis[:, 1:])\n            loss.backward()\n            encoder_optimizer.step()\n            decoder_optimizer.step()\n        encoder.eval()\n        decoder.eval()\n        valid_loss = 0\n        valid_levenshtein = []\n        \n        for images, inchis in tqdm(valid_loader):\n            images = images.to(DEVICE)\n            inchis = inchis.to(DEVICE)\n            with torch.no_grad():\n                encoded = encoder(images)\n                preds = decoder(encoded, inchis)\n                loss = criterion(preds.permute(0, 2, 1), inchis[:, 1:])\n                valid_loss += loss.item()\n                \n                pred_vals = decoder.predict(encoded, max_len = MAX_LEN)\n                vd = levenshtein_score(pred_vals, inchis) #Levenshtein distance\n                valid_levenshtein.append(np.mean(vd))\n                \n        valid_loss /= len(valid_loader)\n        print(f\"[epoch {epoch+1}/{EPOCHS}] loss:{valid_loss}\")\n        valid_levenshtein = sum(valid_levenshtein) / len(valid_levenshtein)\n        print(f\"[epoch {epoch+1}/{EPOCHS}] Lavenshtein Score:{valid_levenshtein}\")\n\n        if valid_loss < best_loss:\n            best_loss = valid_loss\n            torch.save(encoder.state_dict(), \"bms_encoder.pth\")\n            torch.save(decoder.state_dict(), \"bms_decoder.pth\")\n            print(\"Saved...\")\n        else:\n            print(\"Loss did not improve...\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Inference"},{"metadata":{"trusted":true},"cell_type":"code","source":"if PREDICT:\n    encoder = Encoder().to(DEVICE)\n    decoder = Decoder(vocab_size = VOCAB_SIZE, embed_size = 256, encoder_dim = 256, decoder_dim = 512).to(DEVICE)\n    #Must Load weights from train_kernel\n    encoder.load_state_dict(torch.load(\"./bms_encoder.pth\", map_location = DEVICE)) #Change Path!\n    decoder.load_state_dict(torch.load(\"./bms_decoder.pth\", map_location = DEVICE)) #Change Path!\n    \n    preds = []\n    for images in tqdm(test_loader):\n        if time.time() - start_time > time_limit:\n            break #Avoid TimeLimitExceeded Error\n            \n        with torch.no_grad():\n            images = images.to(DEVICE)\n            encoded = encoder(images)\n            pred = decoder.predict(encoded, max_len = MAX_LEN)\n            preds.append(pred)\n    \n    submit_preds = []\n    for pred in preds:\n        submit_preds.append([generate_inchi(p) for p in pred])\n    submit_preds = np.concatenate(submit_preds, axis = 0)\n    print(submit_preds.shape)\n    \n    submit = test[[\"image_id\", \"InChI\"]].copy().reset_index(drop=True)\n    submit.iloc[:len(submit_preds), 1] = submit_preds\n    submit.to_csv(\"submission.csv\", index = False)\n    print(submit.head(10))","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}