{"cells":[{"metadata":{},"cell_type":"markdown","source":"# End to end Tranformer\n\nIn this notebook I extend my previous example on image captioning with transformers (https://www.kaggle.com/juansensio/e2e-transformer-image-captioning-example) for this particular challenge."},{"metadata":{},"cell_type":"markdown","source":"# Data preprocessing\n\nGenerate the vocab, extracted from https://www.kaggle.com/yasufuminakama/inchi-preprocess-2."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nfrom pathlib import Path\n\npath = Path('/kaggle/input/bms-molecular-translation')\nos.listdir(path)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\n\ntrain_labels = pd.read_csv(path / 'train_labels.csv')\ntrain_labels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from tqdm.auto import tqdm\nimport re\ntqdm.pandas()\n\ndef split_form(form):\n    string = ''\n    for i in re.findall(r\"[A-Z][^A-Z]*\", form):\n        elem = re.match(r\"\\D+\", i).group()\n        num = i.replace(elem, \"\")\n        if num == \"\":\n            string += f\"{elem} \"\n        else:\n            string += f\"{elem} {str(num)} \"\n    return string.rstrip(' ')\n\ndef split_form2(form):\n    string = ''\n    for i in re.findall(r\"[a-z][^a-z]*\", form):\n        elem = i[0]\n        num = i.replace(elem, \"\").replace('/', \"\")\n        num_string = ''\n        for j in re.findall(r\"[0-9]+[^0-9]*\", num):\n            num_list = list(re.findall(r'\\d+', j))\n            assert len(num_list) == 1, f\"len(num_list) != 1\"\n            _num = num_list[0]\n            if j == _num:\n                num_string += f\"{_num} \"\n            else:\n                extra = j.replace(_num, \"\")\n                num_string += f\"{_num} {' '.join(list(extra))} \"\n        string += f\"/{elem} {num_string}\"\n    return string.rstrip(' ')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_labels['InChI_1'] = train_labels.InChI.progress_apply(lambda x: x.split('/')[1])\ntrain_labels['InChI_text'] = train_labels['InChI_1'].progress_apply(split_form) + ' ' + train_labels['InChI'].apply(lambda x: '/'.join(x.split('/')[2:])).progress_apply(split_form2).values","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_labels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def compute_vocab(InChIs):\n    special = ['PAD', 'SOS', 'EOS']\n    vocab = special + sorted(list({s for InChI in InChIs for s in InChI}))\n    return vocab","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"VOCAB = compute_vocab(train_labels.InChI_text.map(lambda x: x.split(' ')))\nVOCAB, len(VOCAB)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"lens = train_labels.InChI_text.map(lambda x: x.split(' ')).map(len)\nlens.min(), lens.max()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_labels.to_csv('train_labels_tokenized.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# DataModule"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_image_path(image_id, path=Path('data'), mode=\"train\"):\n    return path / mode / image_id[0] / image_id[1] / image_id[2] / f'{image_id}.png'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom skimage import io\nimport pytorch_lightning as pl\nfrom pathlib import Path\nimport pandas as pd\nfrom sklearn.model_selection import train_test_split\nfrom torch.utils.data import DataLoader\nimport albumentations as A \nimport numpy as np\n\nclass Dataset(torch.utils.data.Dataset):\n    def __init__(self, images, inchis=None, max_len=512, trans=None, train=True, tokens=(0, 1, 2)):\n        self.images = images\n        self.inchis = inchis\n        self.trans = trans\n        self.train = train\n        self.max_len = max_len\n        self.PAD, self.SOS, self.EOS = tokens\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, ix):\n        image = io.imread(self.images[ix]) \n        if self.trans:\n            image = self.trans(image=image)['image']\n        image = torch.tensor(image / 255., dtype=torch.float).unsqueeze(0)\n        if self.train:\n            inchi = torch.tensor([self.SOS] + self.inchis[ix] + [self.EOS], dtype=torch.long)\n            #inchi = torch.nn.functional.pad(inchi, (0, self.max_len - len(inchi)), 'constant', self.PAD)\n            return image, inchi\n        return image\n\n    def collate(self, batch):\n        if self.train:\n            # compute max batch length\n            lens = [len(inchi) for _, inchi in batch]\n            max_len = max(lens)    \n            # pad inchis to max length\n            images, inchis = [], []\n            for image, inchi in batch:\n                images.append(image)\n                inchis.append(torch.nn.functional.pad(inchi, (0, max_len - len(inchi)), 'constant', self.PAD))\n            # optionally, sort by length\n            ixs = torch.argsort(torch.tensor(lens), descending=True)\n            return torch.stack(images)[ixs], torch.stack(inchis)[ixs]\n        return torch.stack([img for img in batch])\n\nclass DataModule(pl.LightningDataModule):\n    def __init__(\n        self, \n        data_file = 'train_labels_tokenized.csv', \n        path=Path('data'), \n        text_column=\"InChI_text\",\n        test_size=0.1, \n        random_state=42, \n        batch_size=64, \n        num_workers=0, \n        pin_memory=True, \n        shuffle_train=True, \n        val_with_train=False,\n        train_trans=None,\n        val_trans=None,\n        subset=None,\n        max_len=512,\n        **kwargs\n    ):\n        super().__init__()\n        self.data_file = data_file\n        self.path = path\n        self.test_size=test_size\n        self.random_state=random_state\n        self.batch_size = batch_size\n        self.num_workers = num_workers\n        self.pin_memory = pin_memory\n        self.shuffle_train = shuffle_train\n        self.val_with_train = val_with_train\n        self.train_trans = train_trans\n        self.val_trans = val_trans\n        self.subset = subset\n        self.max_len = max_len\n        self.text_column = text_column\n        self.stoi = {}\n        self.itos = {}\n\n    def encode(self, InChI):\n        return [self.stoi[token] for token in InChI]\n\n    def decode(self, ixs):\n        skip = [self.stoi['PAD'], self.stoi['SOS'], self.stoi['EOS']]\n        return ('').join([self.itos[ix.item()] for ix in ixs if ix.item() not in skip])\n\n    def setup(self, stage=None):\n        # build indices\n        for i, s in enumerate(VOCAB):\n            self.stoi[s] = i\n        self.itos = {item[1]: item[0] for item in self.stoi.items()}\n        # read csv file with data\n        df = pd.read_csv(self.path / self.data_file)\n        if self.subset:\n            df = df.sample(int(len(df)*self.subset), random_state=self.random_state)\n        # build images paths\n        df.image_id = df.image_id.map(lambda x: get_image_path(x, self.path))\n        # encode inchis\n        df.InChI = df[self.text_column].map(lambda x: x.split(' '))\n        df.InChI = df.InChI.map(self.encode)\n        # train / val splits\n        train, val = train_test_split(df, test_size=self.test_size, random_state=self.random_state, shuffle=True)\n        print(\"Training samples: \", len(train))\n        print(\"Validation samples: \", len(val))\n        # datasets\n        self.train_ds = Dataset(train.image_id.values, train.InChI.values, self.max_len, \n            tokens=(self.stoi['PAD'], self.stoi['SOS'], self.stoi['EOS']), trans = A.Compose([\n            getattr(A, trans)(**params) for trans, params in self.train_trans.items()\n        ]) if self.train_trans else None)\n        self.val_ds = Dataset(val.image_id.values, val.InChI.values, self.max_len, \n            tokens=(self.stoi['PAD'], self.stoi['SOS'], self.stoi['EOS']), trans = A.Compose([\n            getattr(A, trans)(**params) for trans, params in self.val_trans.items()\n        ]) if self.val_trans else None)\n        if self.val_with_train:\n            self.val_ds = self.train_ds           \n    \n    def train_dataloader(self):\n        return DataLoader(\n            self.train_ds, \n            batch_size=self.batch_size, \n            num_workers=self.num_workers, \n            shuffle=self.shuffle_train, \n            pin_memory=self.pin_memory, \n            collate_fn=self.train_ds.collate\n        )\n\n    def val_dataloader(self):\n        return DataLoader(\n            self.val_ds, \n            batch_size=self.batch_size, \n            num_workers=self.num_workers, \n            shuffle=False, \n            pin_memory=self.pin_memory, \n            collate_fn=self.val_ds.collate\n        )\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn.functional as F\nimport pytorch_lightning as pl\nimport numpy as np\nimport torch.nn as nn\n\n# https://github.com/jankrepl/mildlyoverfitted/blob/master/github_adventures/vision_transformer/custom.py\n\nclass PatchEmbedding(nn.Module):\n    def __init__(self, img_size, patch_size, in_chans, embed_dim):\n        super().__init__()\n        self.img_size = img_size\n        self.patch_size = patch_size\n        self.n_patches = (img_size // patch_size) ** 2\n        self.patch_size = patch_size\n        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)\n\n    def forward(self, x):\n        x = self.proj(x)  \n        x = x.flatten(2)  \n        x = x.transpose(1, 2) \n        return x\n\nclass Transformer(pl.LightningModule):\n    def __init__(self, config=None):\n        super().__init__()\n        self.save_hyperparameters(config)\n        self.len_vocab = len(VOCAB)\n        \n        self.patch_embed = PatchEmbedding(self.hparams.img_size, self.hparams.patch_size, 1, self.hparams.embed_dim)\n        self.pos_embed = nn.Parameter(torch.zeros(1, self.patch_embed.n_patches, self.hparams.embed_dim))\n        \n        self.trg_emb = nn.Embedding(self.len_vocab, self.hparams.embed_dim)\n        self.trg_pos_emb = nn.Embedding(self.hparams.max_len, self.hparams.embed_dim)\n\n        dim_feedforward = 4 * self.hparams.embed_dim\n        self.transformer = torch.nn.Transformer(\n            self.hparams.embed_dim, self.hparams.nhead, self.hparams.num_encoder_layers, self.hparams.num_decoder_layers, dim_feedforward, self.hparams.dropout\n        )\n        \n        self.l = nn.LayerNorm(self.hparams.embed_dim)\n        self.fc = nn.Linear(self.hparams.embed_dim, self.len_vocab)\n\n        self.apply(self._init_weights)\n\n    def forward(self, images, captions):\n        # embed images\n        embed_imgs = self.patch_embed(images)\n        embed_imgs = embed_imgs + self.pos_embed  \n        # embed captions\n        B, trg_seq_len = captions.shape \n        trg_positions = (torch.arange(0, trg_seq_len).expand(B, trg_seq_len).to(self.device))\n        embed_trg = self.trg_emb(captions) + self.trg_pos_emb(trg_positions)\n        trg_mask = self.transformer.generate_square_subsequent_mask(trg_seq_len).to(self.device)\n        tgt_padding_mask = captions == 0 # PAD token !!!\n        # transformer\n        y = self.transformer(\n            embed_imgs.permute(1,0,2),  \n            embed_trg.permute(1,0,2),  \n            tgt_mask=trg_mask, \n            tgt_key_padding_mask = tgt_padding_mask\n        ).permute(1,0,2) \n        # head\n        return self.fc(self.l(y))\n\n    # https://github.com/karpathy/minGPT/blob/master/mingpt/model.py\n    \n    def _init_weights(self, module):\n        if isinstance(module, (nn.Linear, nn.Embedding)):\n            module.weight.data.normal_(mean=0.0, std=0.02)\n            if isinstance(module, nn.Linear) and module.bias is not None:\n                module.bias.data.zero_()\n        elif isinstance(module, nn.LayerNorm):\n            module.bias.data.zero_()\n            module.weight.data.fill_(1.0)\n\n    def predict(self, images, SOS=1, EOS=2, temp=1.):\n        self.eval()\n        with torch.no_grad():\n            images = images.to(self.device)\n            B = images.shape[0]\n            # start of sentence\n            trg_input = torch.tensor([SOS], dtype=torch.long, device=self.device).expand(B, 1)\n            while True:\n                # get latest prediction\n                logits = self(images, trg_input)[:,-1,:] / temp\n                probs = F.softmax(logits, dim=-1) \n                # sample\n                pred = torch.multinomial(probs, num_samples=1)\n                # add new prediction\n                trg_input = torch.cat([trg_input, pred], 1)\n                if torch.any(trg_input == EOS, 1).sum().item() == B or trg_input.shape[1] >= self.hparams.max_len:\n                    return trg_input\n\n    def training_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x, y[:,:-1])\n        loss = F.cross_entropy(y_hat.transpose(1,2), y[:,1:]) \n        self.log('loss', loss)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x, y[:,:-1])\n        loss = F.cross_entropy(y_hat.transpose(1,2), y[:,1:]) \n        self.log('val_loss', loss, prog_bar=True)\n    \n    def configure_optimizers(self):\n        optimizer = getattr(torch.optim, self.hparams.optimizer)(self.parameters(), lr=self.hparams.lr)\n        if 'scheduler' in self.hparams:\n            schedulers = [\n                getattr(torch.optim.lr_scheduler, scheduler)(optimizer, **params)\n                for scheduler, params in self.hparams.scheduler.items()\n            ]\n            return [optimizer], schedulers \n        return optimizer","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train"},{"metadata":{"trusted":true},"cell_type":"code","source":"config = {\n    'lr': 0.001,\n    'optimizer': \"Adam\",\n    'batch_size': 64,\n    'gradient_clip_val': 1.0,\n    'num_workers': 4,\n    'pin_memory': True,\n    'subset': 0.1,\n    'img_size': 128,\n    'patch_size': 16,\n    'embed_dim': 64,\n    'nhead': 1,\n    'num_encoder_layers': 1,\n    'num_decoder_layers': 1,\n    'dropout': 0.,\n    'max_len': 277,\n    'train_trans': {\n      'Resize': {\n        'width': 128,\n        'height': 128,\n      }\n    },\n    'val_trans': {\n      'Resize': {\n        'width': 128,\n        'height': 128,\n      }\n    },\n    'gpus': 1,\n    'precision': 16,\n    'max_epochs': 10\n}\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from pytorch_lightning.callbacks import ModelCheckpoint\n\ncheckpoint = ModelCheckpoint(\n    dirpath='./', \n    filename=f'transformer-{{val_loss:.4f}}',\n    save_top_k=1, \n    monitor='val_loss', \n    mode='min'\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dm = DataModule(\n    data_file = '/kaggle/working/train_labels_tokenized.csv', \n    path=Path('/kaggle/input/bms-molecular-translation'), \n    **config\n)\n\nmodel = Transformer(config)\n\ntrainer = pl.Trainer(\n    gpus=config['gpus'],\n    precision=config['precision'],\n    max_epochs=config['max_epochs'],\n    gradient_clip_val=config['gradient_clip_val'],\n    callbacks=[checkpoint]\n)\n\ntrainer.fit(model, dm)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Predictions"},{"metadata":{},"cell_type":"markdown","source":"Compute metric on validation set."},{"metadata":{"trusted":true},"cell_type":"code","source":"from tqdm import tqdm\n\npreds, labels = [], []\nmodel.cuda()\nfor imgs, labs in tqdm(dm.val_dataloader()):\n    outputs = model.predict(imgs)\n    preds += outputs\n    labels += labs.tolist()\n    \nlen(preds)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install python-Levenshtein","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"preds_decoded = [dm.decode(pred) for pred in preds]\npreds_inchis = ['InChI=1S/' + pred for pred in preds_decoded]\n\nlabs_decoded = [dm.decode(lab) for lab in labs]\ninchis = ['InChI=1S/' + lab for lab in labs_decoded]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from Levenshtein import distance\n\nmetric = []\nfor pred, inchi in zip(preds_inchis, inchis):\n    metric.append(distance(pred, inchi))\n    \nnp.mean(metric)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Better than the sample submission !!!"},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_submission = pd.read_csv(path / 'sample_submission.csv')\nsample_submission","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Generate predictions for top 25% test images (the ones used in leaderboard, or are they a random sample?)"},{"metadata":{"trusted":true},"cell_type":"code","source":"limit = int(0.25*len(sample_submission))\nlimit","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_images = sample_submission.image_id[:limit]\ntest_images = test_images.apply(lambda i: get_image_path(i, path, mode=\"test\"))\n\nds = Dataset(test_images, train=False, trans=A.Compose([A.Resize(128,128)]))\n\ndl = torch.utils.data.DataLoader(ds, batch_size=100, num_workers=4, pin_memory=True, shuffle=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# this step takes so long with big models and all test images ...\n\npreds = []\nmodel.cuda()\nfor batch in tqdm(dl):\n    outputs = model.predict(batch)\n    preds += outputs","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"preds_decoded = [dm.decode(pred) for pred in preds]\nsample_submission.InChI[:limit] = ['InChI=1S/'+pred for pred in preds_decoded]\nsample_submission","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_submission.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Where to go from here ?\n\n\nSome ideas:\n\n- Train more epochs\n- Better prediction method (Beam Search?)\n- More layers\n- Data augmentation\n- Higher image resolution\n- Cross Validation and model ensembling\n- Train with all the dataset\n- Generate predictions for all test set\n\nFor me, the biggest stopper is to generate predictions reliably and efficiently. Advice is welcomed :)\n\nMulti-gpu or even TPU training should be relatively straightforward thanks to Pytorch Lightning."}],"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}