{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"from tqdm import tqdm\nimport sys, os\ntry:\n    from transformers import DistilBertModel, DistilBertTokenizer\nexcept:\n    !pip install transformers\n    !pip install sentencepiece\n    from transformers import DistilBertModel, DistilBertTokenizer\nimport torch.nn.functional as F\nimport torch.nn as nn\nimport torch\n\nfrom transformers import AutoModelWithLMHead, AutoTokenizer, AutoModel","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-07-30T10:22:24.803683Z","iopub.execute_input":"2022-07-30T10:22:24.804593Z","iopub.status.idle":"2022-07-30T10:22:28.074128Z","shell.execute_reply.started":"2022-07-30T10:22:24.804485Z","shell.execute_reply":"2022-07-30T10:22:28.072609Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\n\nimport json\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\n#from scipy import sparse\nfrom tqdm import tqdm\nimport os\n\nos.environ[\"WANDB_DISABLED\"] = \"true\"\n\npd.options.display.width = 180\npd.options.display.max_colwidth = 120\n\nBERT_PATH = \"/content/drive/MyDrive/ai4code/models/checkpoint-18000\"\n\ndata_dir = Path('../input/AI4Code')\n\npreprocessed_dir = Path(\"../input/dataset\")","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:22:28.079738Z","iopub.execute_input":"2022-07-30T10:22:28.08072Z","iopub.status.idle":"2022-07-30T10:22:28.089756Z","shell.execute_reply.started":"2022-07-30T10:22:28.080674Z","shell.execute_reply":"2022-07-30T10:22:28.088835Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.read_csv('../input/custom-tokenizer/trainer.csv')","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:22:28.091254Z","iopub.execute_input":"2022-07-30T10:22:28.093028Z","iopub.status.idle":"2022-07-30T10:22:55.879041Z","shell.execute_reply.started":"2022-07-30T10:22:28.092989Z","shell.execute_reply":"2022-07-30T10:22:55.878049Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = df.drop(columns=['Unnamed: 0'])","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:22:55.881757Z","iopub.execute_input":"2022-07-30T10:22:55.882245Z","iopub.status.idle":"2022-07-30T10:22:56.35931Z","shell.execute_reply.started":"2022-07-30T10:22:55.882206Z","shell.execute_reply":"2022-07-30T10:22:56.358171Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#df_copy = df.copy()\n#df = df[df['id'].isin(df['id'].unique()[:5000])]","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:22:56.755575Z","iopub.execute_input":"2022-07-30T10:22:56.756029Z","iopub.status.idle":"2022-07-30T10:22:57.910089Z","shell.execute_reply.started":"2022-07-30T10:22:56.755993Z","shell.execute_reply":"2022-07-30T10:22:57.909091Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = df.set_index('id').set_index('cell_id',append=True)","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:22:57.911473Z","iopub.execute_input":"2022-07-30T10:22:57.911849Z","iopub.status.idle":"2022-07-30T10:22:58.394762Z","shell.execute_reply.started":"2022-07-30T10:22:57.911797Z","shell.execute_reply":"2022-07-30T10:22:58.39373Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 6.2GB->6.7G\n\ndef get_ranks(base, derived):\n    return [base.index(d) for d in derived]\n\ndf_orders = pd.read_csv(\n    data_dir / 'train_orders.csv',\n    index_col='id',\n    squeeze=True,\n).str.split()  # Split the string representation of cell_ids into a list\n\ndf_orders_ = df_orders.to_frame().join(\n    df.reset_index('cell_id').groupby('id')['cell_id'].apply(list),\n    how='right',\n)\n\nranks = {}\nfor id_, cell_order, cell_id in df_orders_.itertuples():\n    ranks[id_] = {'cell_id': cell_id, 'rank': get_ranks(cell_order, cell_id)}\n\ndf_ranks = (\n    pd.DataFrame\n    .from_dict(ranks, orient='index')\n    .rename_axis('id')\n    .apply(pd.Series.explode)\n    .set_index('cell_id', append=True)\n)\n\ndf_ancestors = pd.read_csv(data_dir / 'train_ancestors.csv', index_col='id')\ndf = df.reset_index().merge(df_ranks, on=[\"id\", \"cell_id\"]).merge(df_ancestors[['ancestor_id']], on=[\"id\"])","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:22:58.396269Z","iopub.execute_input":"2022-07-30T10:22:58.396857Z","iopub.status.idle":"2022-07-30T10:23:01.969368Z","shell.execute_reply.started":"2022-07-30T10:22:58.396803Z","shell.execute_reply":"2022-07-30T10:23:01.96822Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:01.970741Z","iopub.execute_input":"2022-07-30T10:23:01.971189Z","iopub.status.idle":"2022-07-30T10:23:01.998667Z","shell.execute_reply.started":"2022-07-30T10:23:01.971147Z","shell.execute_reply":"2022-07-30T10:23:01.997805Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = df.sort_values(['id','cell_type','rank']).reset_index(drop=True)\ndf = df.drop_duplicates(['cell_id'])","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:02.002761Z","iopub.execute_input":"2022-07-30T10:23:02.003044Z","iopub.status.idle":"2022-07-30T10:23:02.22451Z","shell.execute_reply.started":"2022-07-30T10:23:02.00302Z","shell.execute_reply":"2022-07-30T10:23:02.223553Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import gc\nfrom sklearn.model_selection import GroupShuffleSplit\n\nNVALID = 0.00125  # size of validation set\nsplitter = GroupShuffleSplit(n_splits=1, test_size=NVALID, random_state=0)\n\n# Split, keeping notebooks with a common origin (ancestor_id) together\ndf = df.set_index('id')\nids = df.index.unique('id')\nancestors = df_ancestors.loc[ids, 'ancestor_id']\nids_train, ids_valid = next(splitter.split(ids, groups=ancestors))\nids_train, ids_valid = ids[ids_train], ids[ids_valid]\n\ndf_train = df.loc[ids_train, :]\ndf_valid = df.loc[ids_valid, :]\n'''\n# 6GB -> 6.7GB -> 6.3 GB\ndel df_train, df_valid\ngc.collect()\n'''\n\n# 6.3 GB -> 6.3GB \ndel ids, ancestors, ids_train, ids_valid, splitter\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:02.2257Z","iopub.execute_input":"2022-07-30T10:23:02.226083Z","iopub.status.idle":"2022-07-30T10:23:03.476801Z","shell.execute_reply.started":"2022-07-30T10:23:02.226033Z","shell.execute_reply":"2022-07-30T10:23:03.475884Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def generate_triplet(df, mode='train', smoothing = 0.2):\n    triplets = []\n    ids = df.index.unique()\n    random_drop = np.random.random(size=10000)>0.9\n    #random_drop_2 = np.random.random(size=10000)>0.9\n    #random_drop_3 = np.random.random(size=10000)>0.8\n    count = 0\n\n    for id, df_tmp in tqdm(df.groupby('id')):\n        df_tmp_markdown = df_tmp[df_tmp['cell_type']=='markdown']\n\n        df_tmp_code = df_tmp[df_tmp['cell_type']=='code']\n        df_tmp_code_rank = df_tmp_code['rank'].values\n        df_tmp_code_cell_id = df_tmp_code['cell_id'].values\n        df_tmp_code_src_len = df_tmp_code['src_len'].values\n        \n\n        for cell_id, rank, markdown_len in df_tmp_markdown[['cell_id', 'rank', 'src_len']].values:\n            markdown_len = int(markdown_len)\n            #if (markdown_len == 2) and (mode != 'test'):\n            #    continue\n            #labels = np.array([(r==(rank+1)) for r in df_tmp_code_rank]).astype('int')\n            \n            label1 = [(r>(rank)) for r in df_tmp_code_rank]\n            label2 = [(r<(rank)) for r in df_tmp_code_rank]\n            \n            #label3 = np.array([(r==(rank-1)) for r in df_tmp_code_rank]).astype('int')\n            #label4 = np.array([(r==(rank+1)) for r in df_tmp_code_rank]).astype('int')\n            \n            #label5 = label3 ^ label4\n            \n            #labels_down = (np.array(label1[1:] + [True]).astype(int) & np.array(label2).astype(int))\n            \n            labels_up = np.array(label1).astype(int) & np.array([True] + label2[:-1]).astype(int) # rank 크기가 더 큰 code\n\n            \n            #print(\"labels3:\\n\", label3)\n            #label3 = label3 & np.append(1-label4[1:],1)\n            \n            #print(\"labels3:\\n\", label3)\n            # 아래에 markdown이 있을 경우에 제외한다\n            \n            #labels = labels_up\n            labels = labels_up\n            \n            \n            #if len([b for b in label3 if b == 1]) > 0:\n            #    print(\"label3:\\n\", label3)\n            #    print(\"labels:\\n\",labels)\n            \n            non_labels = labels | np.append(np.array([0]), labels[:-1])\n            non_labels = non_labels | np.append(np.array([0]), non_labels[:-1])\n            non_labels = non_labels | np.append(np.array([0]), non_labels[:-1])\n            \n            #print(\"labels:\\n\",labels)\n            #print(\"non_labels:\\n\",non_labels)\n            #time.sleep(1)\n            for cid, code_len, label, non_label in zip(df_tmp_code_cell_id, df_tmp_code_src_len, labels, non_labels):              \n                code_len = int(code_len)\n                count += 1\n                length = min(markdown_len, MAX_MARKDOWN) + code_len\n                length = min(length, MAX_LEN)\n                \n                if mode == 'test':\n                    triplets.append( [cell_id, cid, label, length] )\n                    # triplets.append( [cid, cell_id, label] )\n                elif markdown_len < 1 or code_len < 5:\n                    continue      \n                elif label > 0:\n                    triplets.append( [cell_id, cid, label, length] )\n                elif non_label > 0:\n                    continue\n                    # triplets.append( [cid, cell_id, label] )\n                elif random_drop[count%10000]:\n                    #if code_len == 6:\n                    #    continue\n                    triplets.append( [cell_id, cid, label, length] )\n                    # triplets.append( [cid, cell_id, label] )\n    return triplets","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:03.478136Z","iopub.execute_input":"2022-07-30T10:23:03.478475Z","iopub.status.idle":"2022-07-30T10:23:03.494841Z","shell.execute_reply.started":"2022-07-30T10:23:03.478443Z","shell.execute_reply":"2022-07-30T10:23:03.493786Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def generate_batch_list(triplets, n_bins):\n    batch_list = []\n    start = 0\n    total_length = 0\n    max_length = 0\n    for idx, (_, _, _, length) in enumerate(triplets):\n        if total_length == 0:\n            if int(length) == 128:\n                max_length = int(length) + 1\n            else:\n                max_length = int(length) + 3\n                \n        total_length += max_length\n    \n        if total_length > n_bins:\n            end=idx\n            batch_list.append([start, end, max_length])\n            start = idx+1\n            total_length = 0\n            \n    batch_list.append([start, idx, max_length])\n        \n    return batch_list","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:03.496728Z","iopub.execute_input":"2022-07-30T10:23:03.497304Z","iopub.status.idle":"2022-07-30T10:23:03.511964Z","shell.execute_reply.started":"2022-07-30T10:23:03.497269Z","shell.execute_reply":"2022-07-30T10:23:03.5107Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import DataLoader, Dataset\n\nimport json\ndef str_to_list(document):\n    return list(json.loads(document))\n\nclass MarkdownModel(nn.Module):\n    def __init__(self):\n        super(MarkdownModel, self).__init__()\n        self.distill_bert = AutoModel.from_pretrained(\"../input/pretrained-checkpoint\")\n        self.top = nn.Linear(512, 1)\n        self.dropout = nn.Dropout(0.2)\n        \n    def forward(self, ids, mask):\n        x  = self.distill_bert(ids, mask)[0]\n        #print(x.shape)\n        #logits = x[1]\n        #attentions = x[2]\n        #x = x[0]\n        x = self.dropout(x)\n        x = self.top(x[:, 0, :])\n        x = torch.sigmoid(x) \n        return x #, logits, attentions\n\n\nclass MarkdownDataset(Dataset):\n    \n    def __init__(self, triplets, max_len,max_markdown, dict_cellid_source, mode='train', batch_list=None):\n        super().__init__()\n        self.triplets = triplets\n        self.dict_cellid_source = dict_cellid_source\n        self.max_len = max_len\n        \n        #self.tokenizer = AutoTokenizer.from_pretrained(BERT_PATH, do_lower_case=True)\n        self.tokenizer = AutoTokenizer.from_pretrained(\"../input/custom-tokenizer/my_bpe_tokenizer\", do_lower_case=True)\n        self.tokenizer.add_tokens(['[tsep]'], special_tokens=True)\n        self.mode=mode\n        self.batch_list = batch_list\n        self.max_markdown = max_markdown\n        \n    def __getitem__(self, index):\n        \n        if self.mode == 'train':\n            start = self.batch_list[index][0]\n            end = self.batch_list[index][1]\n            \n            txt = [([2] + str_to_list(self.dict_cellid_source[self.triplets[idx][0]])[:self.max_markdown] + [3]\n                            + str_to_list(self.dict_cellid_source[self.triplets[idx][1]]))[:self.max_len]+ [3]\n                   for idx in range(start, end+1)]\n            labels = torch.FloatTensor([[int(self.triplets[idx][2])] for idx in range(start, end+1)])\n            \n            length = torch.LongTensor([[int(self.dict_cellid_source[self.triplets[idx][1]][1])] for idx in range(start, end+1)])\n            \n            inputs = self.tokenizer.pad({\"input_ids\":txt}, max_length=self.batch_list[index][2], padding=\"max_length\")\n            ids = torch.LongTensor(inputs['input_ids'])\n            mask = torch.LongTensor(inputs['attention_mask'])\n            \n            return ids, mask, labels\n        \n        row = self.triplets[index]\n        txt = ([2] + str_to_list(self.dict_cellid_source[row[0]])[:self.max_markdown] + [3]\n                       + str_to_list(self.dict_cellid_source[row[1]]))[:self.max_len] + [3]\n        \n        label = row[2]\n        \n        inputs = self.tokenizer.pad({\"input_ids\":txt}, max_length=self.max_len+2, padding=\"max_length\")\n        ids = torch.LongTensor(inputs['input_ids'])\n        mask = torch.LongTensor(inputs['attention_mask'])\n\n        return ids, mask, torch.FloatTensor([label]) \n\n    def __len__(self):\n        if self.mode == \"train\":\n            return len(self.batch_list)\n        else:\n            return len(self.triplets)","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:03.513676Z","iopub.execute_input":"2022-07-30T10:23:03.514376Z","iopub.status.idle":"2022-07-30T10:23:03.535543Z","shell.execute_reply.started":"2022-07-30T10:23:03.514341Z","shell.execute_reply":"2022-07-30T10:23:03.534524Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from bisect import bisect\n\ndef count_inversions(a):\n    inversions = 0\n    sorted_so_far = []\n    for i, u in enumerate(a):\n        j = bisect(sorted_so_far, u)\n        inversions += i - j\n        sorted_so_far.insert(j, u)\n    return inversions\n\n\ndef kendall_tau(ground_truth, predictions):\n    total_inversions = 0\n    total_2max = 0  # twice the maximum possible inversions across all instances\n    for gt, pred in zip(ground_truth, predictions):\n        ranks = [gt.index(x) for x in pred]  # rank predicted order in terms of ground truth\n        total_inversions += count_inversions(ranks)\n        n = len(gt)\n        total_2max += n * (n - 1)\n    return 1 - 4 * total_inversions / total_2max","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:03.538573Z","iopub.execute_input":"2022-07-30T10:23:03.53881Z","iopub.status.idle":"2022-07-30T10:23:03.549433Z","shell.execute_reply.started":"2022-07-30T10:23:03.538788Z","shell.execute_reply":"2022-07-30T10:23:03.548393Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import json\ndef str_to_list(document):\n    #print(document)\n    return list(json.loads(document))\n\ndef validate(model, val_loader, mode='train'):\n    model.eval()\n    tbar = tqdm(val_loader, file=sys.stdout)\n    \n    preds = np.zeros(len(val_loader.dataset), dtype='float32')\n    labels = []\n    count = 0\n\n    with torch.no_grad():\n        for idx, data in enumerate(tbar):\n            inputs, target = read_data(data)\n\n            pred = model(inputs[0], inputs[1]).detach().cpu().numpy().ravel()\n\n            preds[count:count+len(pred)] = pred\n            count += len(pred)\n            \n            if mode=='test':\n                labels.append(target.detach().cpu().numpy().ravel())\n            \n    if mode=='test':\n        return preds\n    else:\n        return np.concatenate(labels), np.concatenate(preds)\n\ndef calc_score(data, model, valid_loader):\n    y_valid = validate(model, valid_loader, mode='test')\n    pred_vals = []\n    count = 0\n    for id, df_tmp in tqdm(data.groupby('id'), mininterval=30):\n        df_tmp_mark = df_tmp[df_tmp['cell_type']=='markdown']\n        df_tmp_code = df_tmp[df_tmp['cell_type']!='markdown']\n        df_tmp_code_rank = df_tmp_code['rank'].rank().values\n        N_code = len(df_tmp_code_rank)\n        N_mark = len(df_tmp_mark)\n        preds_tmp = y_valid[count:count+N_mark * N_code]\n        count += N_mark * N_code\n\n        for i in range(N_mark):\n            pred = preds_tmp[i*N_code:i*N_code+N_code] \n            softmax = np.exp((pred-np.mean(pred)) *100-100)/np.sum(np.exp((pred-np.mean(pred)) *100-100))\n            rank = np.sum(softmax * df_tmp_code_rank) - 0.000001\n            pred_vals.append(rank)\n            \n    \n    data.loc[data[\"cell_type\"] == \"markdown\", \"pred\"] = pred_vals\n    pred_df = data.sort_values(\"pred\").groupby(\"id\")[\"cell_id\"].apply(list)\n    true_df = data.sort_values(\"rank\").groupby(\"id\")[\"cell_id\"].apply(list) \n    score = kendall_tau(true_df, pred_df)\n    \n    return score, y_valid, pred_vals","metadata":{"execution":{"iopub.status.busy":"2022-07-30T11:28:55.816792Z","iopub.execute_input":"2022-07-30T11:28:55.817161Z","iopub.status.idle":"2022-07-30T11:28:55.834207Z","shell.execute_reply.started":"2022-07-30T11:28:55.817129Z","shell.execute_reply":"2022-07-30T11:28:55.832713Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def adjust_lr(optimizer, epoch):\n    if epoch < 1:\n        lr = 5e-5\n    elif epoch < 2:\n        lr = 5e-5\n    elif epoch < 5:\n        lr = 5e-5\n    else:\n        lr = 5e-5\n\n    for p in optimizer.param_groups:\n        p['lr'] = lr\n    return lr\n    \ndef get_optimizer(net):\n    param_optimizer = list(net.named_parameters())\n    no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n    optimizer_parameters = [\n        {\n            \"params\": [\n                p for n, p in param_optimizer if not any(nd in n for nd in no_decay)\n            ],\n            \"weight_decay\": 1e-6,\n        },\n        {\n            \"params\": [\n                p for n, p in param_optimizer if any(nd in n for nd in no_decay)\n            ],\n            \"weight_decay\": 0.0,\n        },\n    ]\n    #optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, net.parameters()), lr=1e-5, betas=(0.9, 0.999),\n    #                             eps=1e-8 ) #1e-08)\n    optimizer = torch.optim.AdamW(optimizer_parameters, lr=5e-5)\n    return optimizer\n\nBS = 128+128\nNW = 8\n\n#train_loader = DataLoader(train_ds, batch_size=1, shuffle=True, num_workers=NW,\n#                         pin_memory=False, drop_last=True)","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:03.57041Z","iopub.execute_input":"2022-07-30T10:23:03.570668Z","iopub.status.idle":"2022-07-30T10:23:03.581635Z","shell.execute_reply.started":"2022-07-30T10:23:03.570631Z","shell.execute_reply":"2022-07-30T10:23:03.580765Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def save_checkpoint(model, optimizer, epoch, step):\n    torch.save(\n      {\n      \"model\": model.state_dict(),\n      \"optimizer\": optimizer.state_dict(),\n       #\"scheduler\": scheduler.state_dict()\n      },\n       f\"./checkpoint_{epoch}_{step}.pth\"\n    )","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:03.583888Z","iopub.execute_input":"2022-07-30T10:23:03.584452Z","iopub.status.idle":"2022-07-30T10:23:03.593928Z","shell.execute_reply.started":"2022-07-30T10:23:03.584417Z","shell.execute_reply":"2022-07-30T10:23:03.593005Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 9.5GB -> 10.5GB (훈련시작 시) -> 11GB ~ (훈련) -> 10 GB (강제종료 시)\nimport torch\nimport random\nimport torch.backends.cudnn as cudnn\n\ntorch.manual_seed(0)\ntorch.cuda.manual_seed(0)\ntorch.cuda.manual_seed_all(0)\nnp.random.seed(0)\ncudnn.benchmark = False\ncudnn.deterministic = True\nrandom.seed(0)\n\n\n\ndef read_data(data):\n    return tuple(d.cuda() for d in data[:-1]), data[-1].cuda()\n\ndef read_train_data(data):\n    return tuple(d.cuda() for d in data[:-1]), data[-1].cuda()\n\n\ndef train(model, train_loader, epochs, start, valid_loader, val_df, Type='markdown'):\n    np.random.seed(0)\n    optimizer = get_optimizer(model)\n    mixed_precision = True\n    try:  \n        from apex import amp\n    except:\n        mixed_precision = False  # not installed\n\n    # model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=1)\n    \n    criterion = torch.nn.L1Loss()\n    criterion = torch.nn.BCELoss()\n    \n    for idx in range(start):\n        random.Random(4).shuffle(batch_list)\n\n    for e in range(start, start+epochs):   \n        \n        model.train()\n        random.Random(4).shuffle(batch_list)\n        train_ds.batch_list = batch_list\n        \n        tbar = tqdm(train_ds, file=sys.stdout,mininterval=60)\n        \n        lr = adjust_lr(optimizer, e)\n        \n        #loss_list = []\n        #preds = []\n        #labels = []\n        loss_value = 0\n        \n        for idx, data in enumerate(tbar):\n            inputs, target = read_train_data(data)\n            optimizer.zero_grad()\n            pred = model(inputs[0], inputs[1])\n            loss = criterion(pred, target)\n\n            loss.backward()\n            optimizer.step()\n            \n            loss_value = loss.detach().cpu().item()\n            \n            #avg_loss = np.round(np.mean(loss_list), 4)\n\n            tbar.set_description(f\"Epoch {e+1} Loss: {loss_value} lr: {lr}\")\n            \n            if idx % 5000 == 0:\n                save_checkpoint(model, optimizer, e, idx)\n                score, _, _ = calc_score(val_df, model, valid_loader)\n                print(e, \"epoch\", idx, \"validation score:\", score)\n                with open(\"./log_default_dataset_v3.txt\",'a+') as f:\n                    f.write(str(e) +  \" epoch \" + str(idx) +  \" validation score: \" + str(score))\n\n        save_checkpoint(model, optimizer, e, len(batch_list))\n        score, _, _ = calc_score(val_df, model, valid_loader)\n        print(e+1, \"epoch validation score:\", score)\n\n    return model","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:03.596134Z","iopub.execute_input":"2022-07-30T10:23:03.596405Z","iopub.status.idle":"2022-07-30T10:23:03.615535Z","shell.execute_reply.started":"2022-07-30T10:23:03.596382Z","shell.execute_reply":"2022-07-30T10:23:03.61458Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"n_bins = 50000\nMAX_MARKDOWN = 80\nMAX_LEN = 128\n\nBS = 128 + 128\nNW = 8\n\ndict_cellid_source = dict(zip(df_train['cell_id'].values, df_train['source'].values))\n\ntrain_triplets = generate_triplet(df_train)\n\ntrain_triplets = sorted(train_triplets, key = lambda x: x[3], reverse=True)\n\nbatch_list = generate_batch_list(train_triplets, n_bins)\n\ntrain_ds = MarkdownDataset(train_triplets, max_len=MAX_LEN, max_markdown=MAX_MARKDOWN, dict_cellid_source=dict_cellid_source, mode='train', batch_list=batch_list)\n\ndf_valid[\"pred\"] = df_valid.groupby([\"id\", \"cell_type\"])[\"rank\"].rank()\n\nvalid_triplets = generate_triplet(df_valid, mode='test')\n\nvalid_dict_cellid_source = dict(zip(df_valid['cell_id'].values, df_valid['source'].values))\n\nvalid_ds = MarkdownDataset(valid_triplets, max_len=MAX_LEN, max_markdown=MAX_MARKDOWN, dict_cellid_source=valid_dict_cellid_source, mode='test')\n\nvalid_loader = DataLoader(valid_ds, batch_size=BS*4, shuffle=False, num_workers=NW,\n                          pin_memory=False, drop_last=False)\n\ndf_train_valid = df_train[df_train.index.isin(df_train.index.unique()[:100])].copy()\n\ndf_train_valid[\"pred\"] = df_train_valid.groupby([\"id\", \"cell_type\"])[\"rank\"].rank()\n\ntrain_valid_triplets = generate_triplet(df_train_valid, mode='test')\n\ntrain_valid_dict_cellid_source = dict(zip(df_train_valid['cell_id'].values, df_train_valid['source'].values))\n\ntrain_valid_ds = MarkdownDataset(train_valid_triplets, max_len=MAX_LEN, max_markdown=MAX_MARKDOWN, dict_cellid_source=train_valid_dict_cellid_source, mode='test')\n\ntrain_valid_loader = DataLoader(train_valid_ds, batch_size=BS*4, shuffle=False, num_workers=NW,\n                          pin_memory=False, drop_last=False)\n\nlen(batch_list)","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:03.618815Z","iopub.execute_input":"2022-07-30T10:23:03.619124Z","iopub.status.idle":"2022-07-30T10:23:25.592019Z","shell.execute_reply.started":"2022-07-30T10:23:03.619097Z","shell.execute_reply":"2022-07-30T10:23:25.591048Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch, gc\ngc.collect()\ntorch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:25.595196Z","iopub.execute_input":"2022-07-30T10:23:25.595495Z","iopub.status.idle":"2022-07-30T10:23:26.120472Z","shell.execute_reply.started":"2022-07-30T10:23:25.595469Z","shell.execute_reply":"2022-07-30T10:23:26.119299Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(batch_list)","metadata":{"execution":{"iopub.status.busy":"2022-07-30T10:23:26.122112Z","iopub.execute_input":"2022-07-30T10:23:26.122741Z","iopub.status.idle":"2022-07-30T10:23:26.135615Z","shell.execute_reply.started":"2022-07-30T10:23:26.122697Z","shell.execute_reply":"2022-07-30T10:23:26.134262Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nmodel = MarkdownModel()\n\nmodel = model.cuda()\n\noptimizer = get_optimizer(model)\n\n#model.load_state_dict(torch.load('/content/drive/MyDrive/ai4code/train_model_preprocess_v2/checkpoint_2_20000.pth')['model'])\n#optimizer.load_state_dict(torch.load('/content/drive/MyDrive/ai4code/train_model_preprocess_v2/checkpoint_2_20000.pth')['optimizer'])\n\ninit_epoch = 1\nmodel = train(model, None, epochs=5, start=0,  valid_loader=valid_loader, val_df=df_valid, Type='markdown')","metadata":{"execution":{"iopub.status.busy":"2022-07-30T11:30:00.093465Z","iopub.execute_input":"2022-07-30T11:30:00.09385Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"zeros = [p for _,_,p,_ in train_triplets if p == 0]\nones = [p for _,_,p,_ in train_triplets if p == 1]\nlen(zeros), len(ones)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"score_valid, y_valid,_ = calc_score(df_valid, model, valid_loader)","metadata":{"execution":{"iopub.status.busy":"2022-07-30T11:29:06.885428Z","iopub.execute_input":"2022-07-30T11:29:06.885778Z","iopub.status.idle":"2022-07-30T11:29:36.838006Z","shell.execute_reply.started":"2022-07-30T11:29:06.885748Z","shell.execute_reply":"2022-07-30T11:29:36.837034Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def calc_score_without_valid(data, y_valid, T=20, epsilon=1.1e-6):\n    pred_vals = []\n    count = 0\n    for id, df_tmp in tqdm(data.groupby('id'), mininterval=30):\n        df_tmp_mark = df_tmp[df_tmp['cell_type']=='markdown']\n        df_tmp_code = df_tmp[df_tmp['cell_type']!='markdown']\n        df_tmp_code_rank = df_tmp_code['rank'].rank().values\n        N_code = len(df_tmp_code_rank)\n        N_mark = len(df_tmp_mark)\n        preds_tmp = y_valid[count:count+N_mark * N_code]\n        count += N_mark * N_code\n\n        for i in range(N_mark):\n            pred = preds_tmp[i*N_code:i*N_code+N_code] \n            softmax = np.exp((pred-np.mean(pred)) *100-100)/np.sum(np.exp((pred-np.mean(pred)) *100-100))\n            rank = np.sum(softmax * df_tmp_code_rank) - 0.000001\n            pred_vals.append(rank)\n            \n    \n    data.loc[data[\"cell_type\"] == \"markdown\", \"pred\"] = pred_vals\n    pred_df = data.sort_values(\"pred\").groupby(\"id\")[\"cell_id\"].apply(list)\n    true_df = data.sort_values(\"rank\").groupby(\"id\")[\"cell_id\"].apply(list) \n    score = kendall_tau(true_df, pred_df)\n    return score, y_valid, pred_vals\ncalc_score_without_valid(df_valid, y_valid,2, epsilon=2e-6)[0],","metadata":{"execution":{"iopub.status.busy":"2022-07-30T11:29:36.840349Z","iopub.execute_input":"2022-07-30T11:29:36.84074Z","iopub.status.idle":"2022-07-30T11:29:37.077067Z","shell.execute_reply.started":"2022-07-30T11:29:36.840699Z","shell.execute_reply":"2022-07-30T11:29:37.07601Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"notebook_list = df_valid.index.unique()\n\nscore_list = []\n\nfor nb in notebook_list:\n    pred = df_valid.loc[nb].sort_values(\"pred\").groupby(\"id\")[\"cell_id\"].apply(list)\n    rank = df_valid.loc[nb].sort_values(\"rank\").groupby(\"id\")[\"cell_id\"].apply(list)\n    score = kendall_tau(rank, pred)\n    score_list.append(score)","metadata":{"execution":{"iopub.status.busy":"2022-07-30T11:19:26.293893Z","iopub.execute_input":"2022-07-30T11:19:26.294344Z","iopub.status.idle":"2022-07-30T11:19:26.663615Z","shell.execute_reply.started":"2022-07-30T11:19:26.294304Z","shell.execute_reply":"2022-07-30T11:19:26.658691Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_valid.source = df_valid.source.apply(str_to_list).apply(valid_ds.tokenizer.convert_ids_to_tokens).apply(valid_ds.tokenizer.convert_tokens_to_string)","metadata":{"execution":{"iopub.status.busy":"2022-07-30T11:19:26.692682Z","iopub.execute_input":"2022-07-30T11:19:26.693273Z","iopub.status.idle":"2022-07-30T11:19:27.055931Z","shell.execute_reply.started":"2022-07-30T11:19:26.693237Z","shell.execute_reply":"2022-07-30T11:19:27.052505Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"notebook_list = list(notebook_list)\ndef read_notebook(path):\n    return (\n        pd.read_json(\n            path,\n            dtype={'cell_type': 'category', 'source': 'str'})\n        .assign(id=path.stem)\n        .rename_axis('cell_id')\n    )","metadata":{"execution":{"iopub.status.busy":"2022-07-30T11:19:27.402778Z","iopub.execute_input":"2022-07-30T11:19:27.408803Z","iopub.status.idle":"2022-07-30T11:19:27.415285Z","shell.execute_reply.started":"2022-07-30T11:19:27.408756Z","shell.execute_reply":"2022-07-30T11:19:27.414278Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pos = -1\n\nnb = notebook_list[pos]\n#nb = \"0407d3ec756354\"\n\nbase_dir = Path('../input/AI4Code/train')\n\nprint(nb)\nprint(score_list[pos])\n\nnotebook = read_notebook(base_dir / (nb +\".json\"))\n\nnotebook = notebook.merge(df_valid.loc[nb][['cell_id', 'source', 'pred', 'rank']], on=['cell_id']).drop(columns=['id']).sort_values(['pred'])\n\nnotebook = notebook.reset_index()\n\nnotebook","metadata":{"execution":{"iopub.status.busy":"2022-07-30T11:20:42.808593Z","iopub.execute_input":"2022-07-30T11:20:42.808957Z","iopub.status.idle":"2022-07-30T11:20:42.846906Z","shell.execute_reply.started":"2022-07-30T11:20:42.808925Z","shell.execute_reply":"2022-07-30T11:20:42.845986Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"markdown = notebook[notebook['cell_id'] == '69c66db9'].source_y.values[0]\ncode = notebook[notebook['cell_id'] == '13b0bc44'].source_y.values[0]\n#position = notebook[notebook['cell_id'] == '266763cf'].pct_rank.values[0]\n#code = \"\"\"testing()\"\"\"\n\n#markdown = \"## verbal\"\n\n\n#markdown = preprocess_markdown(markdown)\n#code = preprocess_code(code)\nprint(markdown)\nprint(code)\nmarkdown = [101] + train_ds.tokenizer.encode(markdown, add_special_tokens=False)[:80] + [102]\n#length_info = train_ds.tokenizer.encode(position, add_special_tokens=False)\ncode =  train_ds.tokenizer.encode(code, add_special_tokens=False)[:120] + [102]\n\ntxt = markdown + code\ntxt = txt[:128]\n\nsrc_len = len(txt)\ninputs = train_ds.tokenizer.pad({\"input_ids\":txt})\nids = torch.LongTensor([inputs['input_ids']])\nmask = torch.LongTensor([inputs['attention_mask']])\nlength = torch.LongTensor([[src_len]])\n\nloss = model(ids.cuda(), mask.cuda())\nloss = loss.detach().cpu().numpy().ravel()\n#logits = logits.detach().cpu().numpy().ravel()\n#attentions = [a.detach().cpu().numpy() for a in attentions]\n\nprint(loss)","metadata":{"execution":{"iopub.status.busy":"2022-07-30T11:21:40.630209Z","iopub.execute_input":"2022-07-30T11:21:40.630564Z","iopub.status.idle":"2022-07-30T11:21:40.659782Z","shell.execute_reply.started":"2022-07-30T11:21:40.630535Z","shell.execute_reply":"2022-07-30T11:21:40.658812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"score_list","metadata":{"execution":{"iopub.status.busy":"2022-07-30T11:20:33.27597Z","iopub.execute_input":"2022-07-30T11:20:33.276437Z","iopub.status.idle":"2022-07-30T11:20:33.285988Z","shell.execute_reply.started":"2022-07-30T11:20:33.276401Z","shell.execute_reply":"2022-07-30T11:20:33.284958Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nplt.hist(score_list)","metadata":{"execution":{"iopub.status.busy":"2022-07-30T11:20:17.220604Z","iopub.execute_input":"2022-07-30T11:20:17.221113Z","iopub.status.idle":"2022-07-30T11:20:17.497326Z","shell.execute_reply.started":"2022-07-30T11:20:17.221068Z","shell.execute_reply":"2022-07-30T11:20:17.496148Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}