{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\n\"\"\"\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\"\"\"\n\nos.listdir('/kaggle/input')\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from tqdm.auto import tqdm\ntqdm.pandas()\nimport re","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def 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(' ')\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":"import Levenshtein\nfrom time import time\nimport pickle\n# range more beautiful\nfrom PIL import Image\nimport math\n# nn module\nimport torch\nimport torch.nn as nn\nimport torchvision.models as models\nimport torchvision\nfrom torch.utils.data import Dataset, DataLoader, random_split\nimport torch.optim as optim\nimport torch.nn.functional as F\nimport cv2\n\nos.chdir('/kaggle/input/bms-molecular-translation')\nos.listdir()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def string_to_ints(string,max_len=100):\n    l=[word2num['<BOS>']]\n    if len(string) > max_len:\n        string = string[:max_len]\n    for s in string:\n        l.append(word2num[s])\n    l.append(word2num['<EOS>'])\n    for miss_num in range(max_len - len(string)):\n        l.append(word2num['<PAD>'])\n    return l\ndef string_to_ints_comp(string,max_len=100):\n    l = [word2num['<BOS>']]\n    #string = list(string)\n    if len(string) > max_len:\n        string = string[:max_len]\n    for s in string:\n        l.append(word2num[s])\n    l.append(word2num['<EOS>'])\n    for miss_num in range(max_len - len(string)):\n        l.append(word2num['<PAD>'])\n    return l\ndef ints_to_string(l):\n    l = np.array(l)\n    l = l[l<36]\n    return \"InChI=1S/\"+''.join(list(map(lambda i:num2word[i] ,l)))\ndef ints_to_string_comp(l):\n    l = np.array(l)\n    l = l[l<190]\n    return \"InChI=1S/\"+''.join(list(map(lambda i:num2word[i] ,l)))\ndef idgetpath(idv,type1='train'):\n    #type1 = 'train' or 'test'\n    #path_att = \"/kaggle/input/bms-molecular-translation/\"\n    path = type1+'/{}/{}/{}/{}.png'.format(idv[0],idv[1],idv[2],idv)\n    #print(path)\n    return path\ndef read_pic2(path):\n    f = cv2.imread(path)\n    return inference_transform(f)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"CFG_train = True\nif CFG_train:\n    train = pd.read_csv('/kaggle/input/bms-molecular-translation/train_labels.csv')\n    train['InChI_1'] = train['InChI'].progress_apply(lambda x: x.split('/')[1])\n    train['InChI_text'] = train['InChI_1'].progress_apply(split_form) + ' ' + \\\n                                train['InChI'].apply(lambda x: '/'.join(x.split('/')[2:])).progress_apply(split_form2).values\n    tr_label = train\n    tr_label['filepath'] = tr_label['image_id'].progress_apply(idgetpath)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"datadir = '/kaggle/input/model-data/'\nword2num = pickle.load(open(datadir+'word2num_comp.pkl','rb'))\nnum2word = pickle.load(open(datadir+'num2word_comp.pkl','rb'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torchvision.transforms as transforms\nnorm_mean = [0.485, 0.456, 0.406]\nnorm_std = [0.229, 0.224, 0.225]\ninference_transform = transforms.Compose([\n    transforms.ToPILImage(),\n    transforms.Resize(256),\n    transforms.CenterCrop(224),\n    transforms.ToTensor(),\n    transforms.Normalize(norm_mean, norm_std),\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import math\n\nclass input_pic(Dataset):\n    def __init__(self,path,label,maxlen=160):\n        self.path = path\n        self.loader = read_pic2\n        self.label = label\n        self.maxlen = maxlen\n    def __len__(self):\n        return len(self.path)\n    def __getitem__(self,index):\n        sample = self.loader(self.path[index])\n        label = string_to_ints_comp(self.label[index].split(),self.maxlen)\n        return sample,torch.LongTensor(label),index\nclass EncoderCNN(nn.Module):\n    def __init__(self):\n        super(EncoderCNN, self).__init__()\n        resnet = models.resnet34(pretrained=True)\n        modules = list(resnet.children())[:-2]\n        self.resnet = nn.Sequential(*modules)\n        for param in self.resnet.parameters():\n            param.requires_grad = True\n            \n    def forward(self,x):\n        x = self.resnet(x)\n        x = x.permute(0,2,3,1)#chan is the last it is 512\n        x = x.view(x.size(0),-1,x.size(-1))\n        return x.permute(1,0,2)\n\"\"\"   \nclass EncoderCNN(nn.Module):\n    def __init__(self, embed_size, batch_size=64,res_grad=False,prtrain=True):\n        super(EncoderCNN, self).__init__()\n        resnet = models.resnet18(pretrained=prtrain)\n        for param in resnet.parameters():\n            param.requires_grad_(res_grad)\n        resnet.fc = nn.Linear(resnet.fc.in_features, embed_size*8)\n        self.resnet = resnet\n        self.batch_size = batch_size\n        self.embed_size = embed_size\n\n    def forward(self, image):\n        m = 8\n        features = self.resnet(image)\n        features = features.unsqueeze(1)\n        features = features.view(-1,m,self.embed_size)\n        features = features.permute(1,0,2)\n        return features\n\"\"\"        \nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, dropout=0.1, max_len=100):\n        super(PositionalEncoding, self).__init__()\n        self.dropout = nn.Dropout(p=dropout)\n        #self.len = max_len\n\n        pe = torch.zeros(max_len, d_model)\n        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n        pe[:, 0::2] = torch.sin(position * div_term)\n        pe[:, 1::2] = torch.cos(position * div_term)\n        pe = pe.unsqueeze(0).transpose(0, 1)\n        self.register_buffer('pe', pe)\n\n    def forward(self, x):\n        x = x + self.pe[:x.size(0), :]\n        return self.dropout(x)\nclass TransformerDecoderLayer_pre(nn.TransformerDecoderLayer):\n    def __init__(self, d_model, nhead):\n        super().__init__(d_model=d_model, nhead=nhead)\n\n    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,\n                tgt_key_padding_mask=None, memory_key_padding_mask=None):\n        tgt_ = self.norm1(tgt)\n        tgt2 = self.self_attn(tgt_, tgt_, tgt_, attn_mask=tgt_mask,\n                              key_padding_mask=tgt_key_padding_mask)[0]\n        tgt = tgt + self.dropout1(tgt2)\n        tgt_ = self.norm2(tgt)\n        tgt2 = self.multihead_attn(tgt_, memory, memory, attn_mask=memory_mask,\n                                   key_padding_mask=memory_key_padding_mask)[0]\n        tgt = tgt + self.dropout2(tgt2)\n        tgt_ = self.norm3(tgt)\n        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt_))))\n        tgt = tgt + self.dropout3(tgt2)\n        #tgt = self.norm3(tgt)\n        return tgt\n\n\nclass Decoder(nn.Module):\n    def __init__(self,d_model=512,n_head=8,num_layer=6,outtoken=193,dropout=0.1,batch_size=64,length=160):\n        super(Decoder, self).__init__()\n        self.decoder_layer = TransformerDecoderLayer_pre(d_model=d_model, nhead=n_head)\n        self.transformdecoder = nn.TransformerDecoder(self.decoder_layer, num_layers=num_layer)\n        self.fc_out = nn.Linear(d_model, outtoken)\n        self.decoder = nn.Embedding(outtoken, d_model)\n        self.pos_decoder = PositionalEncoding(d_model, dropout,max_len=length)\n        \n    def forward(self,memory,tgc,tgt_mask=None):\n        tgc = self.decoder(tgc)\n        #print(tgc.size())\n        tgc = self.pos_decoder(tgc)\n        #memory = self.resn(src1,src2,src3,src4)\n        #print(memory.shape)\n        #print(tgc.shape)\n        if tgt_mask is None:\n            tgt_mask = self.generate_square_subsequent_mask(len(tgc)).to(tgc.device)\n        if memory.size()[-1] == tgc.size()[-1]:\n            output = self.transformdecoder(tgc,memory,tgt_mask=tgt_mask)\n            output = self.fc_out(output)\n        else:\n            print(\"src and tgc dim different!\")\n        return output\n        \n    def generate_square_subsequent_mask(self, sz):\n        r\"\"\"Generate a square mask for the sequence. The masked positions are filled with float('-inf').\n            Unmasked positions are filled with float(0.0).\n        \"\"\"\n        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)\n        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n        return mask\n    def predict(self,memory,tgc=None,aim_length=160):\n        \"\"\" \n        the memory is 85 X 1 x d_model,output is list\n        \"\"\"\n        epo = memory.size()[1]\n        pred_num = torch.full((1,epo),word2num['<BOS>']).cuda()\n        #pred = self.decoder(pred_num)\n        pred_list = torch.zeros((epo,aim_length),dtype=torch.long).cuda()\n        pred_list[:,0] = pred_num.reshape(-1)\n        for i in range(1,aim_length):\n            pred = self.decoder(pred_list[:,:i].transpose(0,1))\n            #print(pred.shape)\n            pred = self.pos_decoder(pred)\n            #print(pred.shape)\n            output = self.transformdecoder(pred,memory)\n            output = self.fc_out(output)\n            #print(output.size())\n            #return output,pred_list\n            output = output[-1].argmax(1).reshape(-1)\n            #print(output.shape)\n            pred_list[:,i] = output\n        return pred_list\nclass Resnet_transdecoder(nn.Module):\n    def __init__(self,d_model=512,num_layer=3,batch_size=64,length=160):\n        super(Resnet_transdecoder,self).__init__()\n        self.resn = EncoderCNN()\n        self.decoder = Decoder(d_model=d_model,num_layer=num_layer,batch_size=batch_size,length=length)\n    def forward(self,src1,tgc):\n        memory = self.resn(src1)\n        output = self.decoder(memory,tgc)\n        return output","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class LabelSmoothCELoss(nn.Module):\n    def __init__(self):\n        super().__init__()\n\n    def forward(self, pred, label, smoothing=0.1):\n        pred = F.softmax(pred, dim=1)\n        one_hot_label = F.one_hot(label, pred.size(1)).float()\n        smoothed_one_hot_label = (\n            1.0 - smoothing) * one_hot_label + smoothing / pred.size(1)\n        loss = (-torch.log(pred)) * smoothed_one_hot_label\n        loss = loss.sum(axis=1, keepdim=False)\n        loss = loss.mean()\n\n        return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from time import ctime\nembedding_dim = 512\n#vocab_len = 39 #39\nvocab_len=193\nepoches = 3\ninput_len = 160 # adjust small?\nsave_time = 1\nprint_time = 1000\ndef train(train_sample,model_path=None,model_dir='/content/',d_model=512,epoches=3,accumulation=8,batch_size=64,output_token=193,print_time=1000):\n    testda = input_pic(train_sample['filepath'],train_sample['InChI_text'],maxlen=158)\n\n    train_len = int(len(testda) * 0.9)\n    val_len = len(testda) - train_len\n    train_set, val_set = random_split(testda, [train_len, val_len])\n\n    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=4)\n    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=4)\n    #testdata = DataLoader(testda,batch_size=batch_size,shuffle=False,num_workers=1)\n    #resn = EncoderCNN(embed_size=d_model,batch_size=batch_size)\n    model = Resnet_transdecoder(d_model=d_model,num_layer=1,batch_size=64,length=160)\n    if model_path:\n        model.load_state_dict(torch.load(model_path))\n    print(\"train length is {}, val length is {}\".format(len(train_loader),len(val_loader)))\n    #resn = resn.cuda()\n    model = model.cuda()\n    # optimizer = optim.SGD(model.parameters(), lr=0.5)\n    optimizer = optim.Adam(model.parameters(),lr=5e-3,weight_decay=1e-6)\n    #scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)\n    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max = 4, eta_min = 1e-6, last_epoch = -1 )\n    #criterion = nn.CrossEntropyLoss()\n    criterion = LabelSmoothCELoss()\n    best_loss = 100\n    for i in range(epoches):\n        model.train()\n        epoch_loss = 0\n        for index,(pic,lab,ind) in enumerate(train_loader):\n            pic = pic.cuda()\n            lab = lab.transpose(0,1).cuda()\n            output = model(pic,lab[:-1,:])\n            n = output.shape[-1]\n            loss = criterion(output.reshape(-1, n), lab[1:, :].reshape(-1))\n            loss = loss / accumulation # 累计数量梯度更新\n            loss.backward()\n            nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            if (index+1) % accumulation == 0:\n                optimizer.step()\n                optimizer.zero_grad()\n            epoch_loss += loss.item()\n            if index % print_time == 0:\n                print(\"current error is {}\".format(loss.item()))\n                print(\"current time is {}\",end=\" \")\n                print(ctime())\n        epoch_loss_tra = epoch_loss/len(train_loader)\n        print(\"train loss aveg batch is {}\".format(epoch_loss_tra))\n        model.eval()\n        epoch_loss = 0\n        with torch.no_grad():\n            valid_dis = []\n            for i, batch in enumerate(val_loader):\n                pic,lab,ind = batch\n                pic = pic.cuda()\n                lab = lab.transpose(0,1).cuda()\n                memory = model.resn(pic)\n                #output = model(pic,lab[:-1,:])\n                output = model.decoder(memory,lab[:-1,:])\n                n = output.shape[-1]\n                loss = criterion(output.reshape(-1, n), lab[1:, :].reshape(-1))\n                epoch_loss += loss.item()\n                textlist = model.decoder.predict(memory,aim_length=160)\n                textlist = textlist.cpu().detach().numpy()\n                texts = [ints_to_string_comp(text) for text in textlist]\n                lab = lab.transpose(0,1).cpu().detach().numpy()\n                truetext = [ints_to_string_comp(text) for text in lab]\n                dis = []\n                for row in range(len(texts)):\n                    dis.append(Levenshtein.distance(texts[row],truetext[row]))\n                valid_dis.append(np.mean(dis))\n            epoch_loss_val = epoch_loss / len(val_loader)\n            print('valid data loss is {}'.format(epoch_loss_val))\n            print('valid data Levensh dist is {}'.format(np.mean(valid_dis)))\n            if epoch_loss_val < best_loss:\n                best_loss = epoch_loss_val\n                model_name = model_dir+\"model_{0:.5f}.pt\".format(epoch_loss_val)\n                print(model_name)\n                torch.save(model.state_dict(), model_name)\n        print('save finished')\n    return model_name","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#import pynvml\ndef predict(model,idv,length = 160,epo=64,types='test'):\n    n = len(idv)\n    model = model.cuda()\n    model.eval()\n    rs_str = []\n    print('test data num :{}'.format(n))\n    #pynvml.nvmlInit()\n    #handle = pynvml.nvmlDeviceGetHandleByIndex(0)\n    #meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)\n    with torch.no_grad():\n        tp_list = []\n        epo_num = int(n/epo)\n        print('epoches num is {}'.format(epo_num))\n        for i in range(epo_num+1):\n            pic = torch.zeros(epo,3,224,224)\n            if (i+1) * epo > n:\n                for j , num_i in enumerate(range(i*epo,n)):\n                    pic_id = read_pic2(idgetpath(idv[num_i],type1=types))\n                    pic[j] = pic_id\n                epo_act = j + 1\n            else:\n                for j,num_i in enumerate(range(i*epo,i*epo+epo)):\n                    pic_id = read_pic2(idgetpath(idv[num_i],type1=types))\n                    pic[j] = pic_id\n                epo_act = epo\n                print(\"for fin\")\n            if i == 0:\n                print('success data')\n            pic = pic[:epo_act].cuda()\n            pic = pic.to(torch.float32)\n            memory = model.resn(pic)\n            #print(memory.size())\n            #print(memory.shape)\n            textlist = model.decoder.predict(memory,aim_length=length)#.transpose(0,1)\n            if i == 0:\n                print(\"success output\")\n            #pred_num = [word2num['<BOS>']]\n            #pre = prior\n            #print(textlist.shape)\n            if i % 20 == 0:\n                print(\"1 w data predict!\")\n            textlist = textlist.cpu().detach().numpy()\n            texts = [ints_to_string_comp(text) for text in textlist]\n            tp_list = tp_list + texts\n    return tp_list","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if CFG_train:\n    st = time()\n    #model_path = '/kaggle/input/model-deal/model_0.38631.pt'\n    mdname = train(tr_label[:1000000],model_dir='/kaggle/working/',epoches=3)\n    print(time()-st)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not CFG_train:\n    st = time()\n    model = Resnet_transdecoder(d_model=512,batch_size=512,length=160)\n    #model.load_state_dict(torch.load(mdname))\n    test = pd.read_csv('sample_submission.csv')\n    testeg = test[:10000]\n    testdata = predict(model,testeg['image_id'],length=160,epo=512,types='test')\n    x = pd.DataFrame({'image_id':testeg.image_id,'InChI':testdata})\n    x.to_csv(datadir+'partdata1.csv')\n    print(time()-st)","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}