{"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":"import os\nfrom pathlib import Path\nimport pandas as pd\nfrom tqdm import tqdm\nimport numpy as np\n\nclass Preprocessor:\n    def __init__(self):\n        self.data_dir = Path('/kaggle/input/AI4Code')\n    \n    def read_notebook(self, path):\n        return (\n            pd.read_json(path, dtype={'cell_type': 'category', 'source': 'str'})\n                .assign(id=path.stem)\n                .rename_axis('cell_id')\n            )\n    \n    def get_notebooks_test(self):\n        paths_train = list((self.data_dir / 'test').glob('*.json'))\n        notebooks_test = [\n            self.read_notebook(path) for path in tqdm(paths_train, desc='Test NBs')\n        ]\n        return notebooks_test\n    \n    def get_test_df(self):\n        notebooks_test = self.get_notebooks_test()\n        df = (\n            pd.concat(notebooks_test)\n                .set_index('id', append=True)\n                .swaplevel()\n                .sort_index(level='id', sort_remaining=False)\n        )\n        return df\n    \n    def clean_code(self, cell):\n        return str(cell).replace('\\\\n', '\\n')\n    \n    def sample_cells(self, cells, n=20):\n        cells = [self.clean_code(cell) for cell in cells]\n\n        if n >= len(cells): # 코드 셀이 20개 이하라면 그냥 반환\n            return [cell[:200] for cell in cells]\n        else:\n            results = []\n            step = len(cells) / n # 총 20개의 코드셀이 샘플링 되도록 스텝을 조절\n            idx = 0\n            while int(np.round(idx) < len(cells)):\n                results.append(cells[int(np.round(idx))])\n                idx += step\n            assert cells[0] in results # 첫번쨰 코드셀은 반드시 들어가야 한다?\n            if cells[-1] not in results: # 말전 코드셀은 반드시 들어가야 한다?\n                results[-1] = cells[-1]\n            return results\n        \n    def get_features(self, df):\n        features = dict()\n        df = df.sort_values('rank').reset_index(drop=True)\n\n        for idx, sub_df in tqdm(df.groupby('id')):\n            features[idx] = dict()\n            total_md = sub_df[sub_df.cell_type == 'markdown'].shape[0]\n            code_sub_df = sub_df[sub_df.cell_type == 'code']\n            total_code = code_sub_df.shape[0]\n            codes = self.sample_cells(code_sub_df.source.values, 20)\n            features[idx]['total_code'] = total_code\n            features[idx]['total_md'] = total_md\n            features[idx]['codes'] = codes\n        \n        return features","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-08-01T03:39:24.648169Z","iopub.execute_input":"2022-08-01T03:39:24.648529Z","iopub.status.idle":"2022-08-01T03:39:24.666230Z","shell.execute_reply.started":"2022-08-01T03:39:24.648498Z","shell.execute_reply":"2022-08-01T03:39:24.665196Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"preprocessor = Preprocessor()\n\ntest_df = preprocessor.get_test_df().reset_index()\ntest_df['rank'] = test_df.groupby(['id', 'cell_type']).cumcount()\ntest_df['pred'] = test_df.groupby(['id', 'cell_type'])['rank'].rank(pct=True)\n\ntest_fts = preprocessor.get_features(test_df)","metadata":{"execution":{"iopub.status.busy":"2022-08-01T03:39:24.668169Z","iopub.execute_input":"2022-08-01T03:39:24.668766Z","iopub.status.idle":"2022-08-01T03:39:24.731413Z","shell.execute_reply.started":"2022-08-01T03:39:24.668726Z","shell.execute_reply":"2022-08-01T03:39:24.730391Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import DataLoader, Dataset\nimport torch\nimport torch.nn as nn\nfrom transformers import AutoTokenizer, AutoModel\nimport sys\n\nclass MarkdownModel(nn.Module):\n    def __init__(self, model_path):\n        super(MarkdownModel, self).__init__()\n\n        self.model = AutoModel.from_pretrained(model_path)\n        self.top = nn.Linear(769, 1)\n\n    def forward(self, ids, mask, fts):\n        x = self.model(ids, mask)[0]\n        x = torch.cat((x[:, 0, :], fts), 1)\n        x = self.top(x)\n        return x\n\nclass MarkdownDataset(Dataset):\n    def __init__(self, df, model_name_or_path, total_max_len, md_max_len, fts):\n        super().__init__()\n        \n        self.df = df.reset_index(drop=True)\n        self.md_max_len = md_max_len\n        self.total_max_len = total_max_len\n        self.fts = fts\n        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n\n    def __getitem__(self, index):\n        row = self.df.iloc[index]\n\n        inputs = self.tokenizer.encode_plus(\n            row.source,\n            None,\n            add_special_tokens=True,\n            max_length=self.md_max_len,\n            padding='max_length',\n            return_token_type_ids=True,\n            truncation=True\n        )\n\n        code_inputs = self.tokenizer.batch_encode_plus(\n            [str(x) for x in self.fts[row.id]['codes']],\n            add_special_tokens=True,\n            max_length=23,\n            padding='max_length',\n            truncation=True\n        )\n\n        n_md = self.fts[row.id]['total_md']\n        n_code = self.fts[row.id]['total_code']\n        if n_md + n_code == 0:\n            fts = torch.FloatTensor([0])\n        else:\n            fts = torch.FloatTensor([n_md / (n_md + n_code)])\n\n        ids = inputs['input_ids']\n        for x in code_inputs['input_ids']:\n            ids.extend(x[:-1])\n        ids = ids[:self.total_max_len]\n        if len(ids) != self.total_max_len:\n            ids = ids + [self.tokenizer.pad_token_id, ] * (self.total_max_len - len(ids))\n        ids = torch.LongTensor(ids)\n\n        mask = inputs['attention_mask']\n        for x in code_inputs['attention_mask']:\n            mask.extend(x[:-1])\n        mask = mask[:self.total_max_len]\n        if len(mask) != self.total_max_len:\n            mask = mask + [self.tokenizer.pad_token_id, ] * (self.total_max_len - len(mask))\n        mask = torch.LongTensor(mask)\n\n        assert len(ids) == len(mask)\n\n        return ids, mask, fts, torch.FloatTensor([row.pct_rank])\n\n    def __len__(self):\n        return self.df.shape[0]            \n\ndef read_data(data):\n    return tuple(d.cuda() for d in data[:-1]), data[-1].cuda()\n\ndef validate(model, val_loader):\n    model.eval()\n    \n    tbar = tqdm(val_loader, file=sys.stdout)\n    \n    preds = []\n    labels = []\n\n    with torch.no_grad():\n        for idx, data in enumerate(tbar):\n            inputs, target = read_data(data)\n\n            pred = model(*inputs)\n\n            preds.append(pred.detach().cpu().numpy().ravel())\n            labels.append(target.detach().cpu().numpy().ravel())\n    \n    return np.concatenate(labels), np.concatenate(preds)\n\ndef predict(model_path, ckpt_path):\n    model = MarkdownModel(model_path)\n    model = model.cuda()\n    model.eval()\n    model.load_state_dict(torch.load(ckpt_path))\n    BS = 32\n    NW = 2\n    MAX_LEN = 64\n    test_df['pct_rank'] = 0\n    test_ds = MarkdownDataset(test_df[test_df['cell_type'] == 'markdown'].reset_index(drop=True),\n                              md_max_len=64,\n                              total_max_len=512,\n                              model_name_or_path=model_path,\n                              fts=test_fts)\n    test_loader = DataLoader(test_ds, batch_size=BS, shuffle=False, num_workers=NW,\n                             pin_memory=False, drop_last=False)\n    _, y_test = validate(model, test_loader)\n    return y_test","metadata":{"execution":{"iopub.status.busy":"2022-08-01T03:39:24.733731Z","iopub.execute_input":"2022-08-01T03:39:24.734351Z","iopub.status.idle":"2022-08-01T03:39:24.758163Z","shell.execute_reply.started":"2022-08-01T03:39:24.734314Z","shell.execute_reply":"2022-08-01T03:39:24.757271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_path = \"/kaggle/input/huggingface-code-models/graphcodebert-base\"\nckpt_path = \"/kaggle/input/codebert2/model_epoch_5_0.8499189849500974.bin\"\ny_test = predict(model_path, ckpt_path)","metadata":{"execution":{"iopub.status.busy":"2022-08-01T03:39:24.760144Z","iopub.execute_input":"2022-08-01T03:39:24.760637Z","iopub.status.idle":"2022-08-01T03:39:34.018155Z","shell.execute_reply.started":"2022-08-01T03:39:24.760597Z","shell.execute_reply":"2022-08-01T03:39:34.016555Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df.loc[test_df[\"cell_type\"] == \"markdown\", \"pred\"] = y_test\nsub_df = test_df.sort_values(\"pred\").groupby(\"id\")[\"cell_id\"].apply(lambda x: \" \".join(x)).reset_index()\nsub_df.rename(columns={\"cell_id\": \"cell_order\"}, inplace=True)\nsub_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-08-01T03:39:34.021296Z","iopub.execute_input":"2022-08-01T03:39:34.021875Z","iopub.status.idle":"2022-08-01T03:39:34.045041Z","shell.execute_reply.started":"2022-08-01T03:39:34.021842Z","shell.execute_reply":"2022-08-01T03:39:34.044086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub_df.to_csv(\"submission1.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2022-08-01T03:39:34.046556Z","iopub.execute_input":"2022-08-01T03:39:34.046929Z","iopub.status.idle":"2022-08-01T03:39:34.054239Z","shell.execute_reply.started":"2022-08-01T03:39:34.046893Z","shell.execute_reply":"2022-08-01T03:39:34.052620Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Pairwise: submission2.csv","metadata":{}},{"cell_type":"code","source":"import joblib\nimport re\nimport nltk\nfrom nltk.stem import WordNetLemmatizer\nnltk.download('wordnet')\nstemmer = WordNetLemmatizer()\n\nPRETRAINED_MODEL_PATH = '/kaggle/input/pairwisemodels/checkpoint-18000'\nFINETUNED_MODEL_PATH = '/kaggle/input/pairwisemodels//my_own_model.bin'\nTOKENIZER_PATH = '/kaggle/input/pairwisemodels//my_own_tokenizer'\n\ndata_dir = Path('/kaggle/input/AI4Code')\n\ndef generate_triplet(df, mode='train'):\n  triplets = []\n  ids = df.id.unique()\n  random_drop = np.random.random(size=10000)>0.9\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\n    for cell_id, rank in df_tmp_markdown[['cell_id', 'rank']].values:\n      labels = np.array([(r==(rank+1)) for r in df_tmp_code_rank]).astype('int')\n\n      for cid, label in zip(df_tmp_code_cell_id, labels):\n        count += 1\n        if label==1:\n          triplets.append( [cell_id, cid, label] )\n          # triplets.append( [cid, cell_id, label] )\n        elif mode == 'test':\n          triplets.append( [cell_id, cid, label] )\n          # triplets.append( [cid, cell_id, label] )\n        elif random_drop[count%10000]:\n          triplets.append( [cell_id, cid, label] )\n          # triplets.append( [cid, cell_id, label] )\n    \n  return triplets\n\ndef preprocess_text(document):\n    # Remove all the special characters\n    document = re.sub(r'\\W', ' ', str(document))\n\n    # remove all single characters\n    document = re.sub(r'\\s+[a-zA-Z]\\s+', ' ', document)\n    \n    # Remove single characters from the start\n    document = re.sub(r'\\^[a-zA-Z]\\s+', ' ', document)\n    \n    # Substituting multiple spaces with single space\n    document = re.sub(r'\\s+', ' ', document, flags=re.I)\n    \n    # Removing prefixed 'b'\n    document = re.sub(r'^b\\s+', '', document)\n    \n    # Converting to Lowercase\n    document = document.lower()\n    #return document\n    \n    # Lemmatization\n    tokens = document.split()\n    tokens = [stemmer.lemmatize(word) for word in tokens]\n    tokens = [word for word in tokens if len(word) > 3]\n    preprocessed_text = ' '.join(tokens)\n    return preprocessed_text\n\nclass MarkdownModel(nn.Module):\n    def __init__(self):\n        super(MarkdownModel, self).__init__()\n        self.distill_bert = AutoModel.from_pretrained(PRETRAINED_MODEL_PATH)\n        self.top = nn.Linear(512, 1)\n\n        self.dropout = nn.Dropout(0.2)\n        \n    def forward(self, ids, mask):\n        x = self.distill_bert(ids, mask)[0]\n        x = self.dropout(x)\n        x = self.top(x[:, 0, :])\n        x = torch.sigmoid(x) \n        return x\n\nclass MarkdownDataset(Dataset):\n    \n    def __init__(self, df, max_len, mode='train'):\n        super().__init__()\n        self.df = df\n        self.max_len = max_len\n        self.tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH, do_lower_case=True)\n        self.mode=mode\n\n    def __getitem__(self, index):\n        row = self.df[index]\n\n        label = row[-1]\n\n        txt = dict_cellid_source[row[0]] + '[SEP]' + dict_cellid_source[row[1]]\n\n        inputs = self.tokenizer.encode_plus(\n            txt,\n            None,\n            add_special_tokens=True,\n            max_length=self.max_len,\n            padding=\"max_length\",\n            return_token_type_ids=True,\n            truncation=True\n        )\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        return len(self.df)\n\ndef read_data(data):\n    return tuple(d.cuda() for d in data[:-1]), data[-1].cuda()\n\n\ndef validate(model, val_loader, mode='train'):\n    model.eval()\n    \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    if mode=='test':\n      return preds\n    else:\n      return np.concatenate(labels), np.concatenate(preds)\n\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-08-01T03:39:34.058124Z","iopub.execute_input":"2022-08-01T03:39:34.058382Z","iopub.status.idle":"2022-08-01T03:39:54.120384Z","shell.execute_reply.started":"2022-08-01T03:39:34.058359Z","shell.execute_reply":"2022-08-01T03:39:54.119292Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BS = 128\nNW = 8\nMAX_LEN = 128\n\npaths_test = list((data_dir / 'test').glob('*.json'))\nnotebooks_test = [\n    read_notebook(path) for path in tqdm(paths_test, desc='Test NBs')\n]\ntest_df = (\n    pd.concat(notebooks_test)\n    .set_index('id', append=True)\n    .swaplevel()\n    .sort_index(level='id', sort_remaining=False)\n).reset_index()\n\ntest_df.source = test_df.source.apply(preprocess_text)\ndict_cellid_source = dict(zip(test_df['cell_id'].values, test_df['source'].values))\ntest_df[\"rank\"] = test_df.groupby([\"id\", \"cell_type\"]).cumcount()\ntest_df[\"pred\"] = test_df.groupby([\"id\", \"cell_type\"])[\"rank\"].rank(pct=False)\ntest_triplets = generate_triplet(test_df, mode = 'test')\ntest_df[\"pct_rank\"] = 0\ntest_ds = MarkdownDataset(test_triplets, max_len=MAX_LEN)\ntest_loader = DataLoader(test_ds, batch_size=BS * 4, shuffle=False, num_workers=NW,\n                          pin_memory=False, drop_last=False)\nimport gc \ngc.collect()\nlen(test_ds), test_ds[0]","metadata":{"execution":{"iopub.status.busy":"2022-08-01T03:39:54.121815Z","iopub.execute_input":"2022-08-01T03:39:54.122257Z","iopub.status.idle":"2022-08-01T03:39:54.475506Z","shell.execute_reply.started":"2022-08-01T03:39:54.122216Z","shell.execute_reply":"2022-08-01T03:39:54.474615Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = MarkdownModel()\nmodel = model.cuda()\nmodel.load_state_dict(torch.load(FINETUNED_MODEL_PATH))\ny_test = validate(model, test_loader, mode='test')","metadata":{"execution":{"iopub.status.busy":"2022-08-01T03:39:54.476998Z","iopub.execute_input":"2022-08-01T03:39:54.477301Z","iopub.status.idle":"2022-08-01T03:39:56.320109Z","shell.execute_reply.started":"2022-08-01T03:39:54.477276Z","shell.execute_reply":"2022-08-01T03:39:56.318870Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"preds_copy = y_test\npred_vals = []\ncount = 0\nfor id, df_tmp in tqdm(test_df.groupby('id')):\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    \n    preds_tmp = preds_copy[count:count+N_mark * N_code]\n    \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    \n      softmax = np.exp((pred-np.mean(pred)) *20)/np.sum(np.exp((pred-np.mean(pred)) *20)) \n    \n      rank = np.sum(softmax * df_tmp_code_rank)\n      pred_vals.append(rank)\n    \ndel model\ndel test_triplets[:]\ndel dict_cellid_source\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2022-08-01T03:39:56.321746Z","iopub.execute_input":"2022-08-01T03:39:56.322809Z","iopub.status.idle":"2022-08-01T03:39:56.560513Z","shell.execute_reply.started":"2022-08-01T03:39:56.322766Z","shell.execute_reply":"2022-08-01T03:39:56.559622Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df.loc[test_df[\"cell_type\"] == \"markdown\", \"pred\"] = pred_vals\nsub_df = test_df.sort_values(\"pred\").groupby(\"id\")[\"cell_id\"].apply(lambda x: \" \".join(x)).reset_index()\nsub_df.rename(columns={\"cell_id\": \"cell_order\"}, inplace=True)\nsub_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-08-01T03:39:56.564231Z","iopub.execute_input":"2022-08-01T03:39:56.565073Z","iopub.status.idle":"2022-08-01T03:39:56.580618Z","shell.execute_reply.started":"2022-08-01T03:39:56.565035Z","shell.execute_reply":"2022-08-01T03:39:56.579576Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub_df.to_csv(\"submission2.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2022-08-01T03:39:56.582235Z","iopub.execute_input":"2022-08-01T03:39:56.582792Z","iopub.status.idle":"2022-08-01T03:39:56.590226Z","shell.execute_reply.started":"2022-08-01T03:39:56.582751Z","shell.execute_reply":"2022-08-01T03:39:56.589150Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Rank Ensemble","metadata":{}},{"cell_type":"code","source":"# Reading the submissions\ndf_1 = pd.read_csv('submission1.csv')\ndf_2 = pd.read_csv('submission2.csv')\n\n# Averaging the indices and sorting the resulting submission by the aggregated ensembled indices\nnew_samples = []\nfor sample_idx in range(len(df_1)):\n    # {'0a226b6a': 0, ...}\n    sample_1 = {k: v for v, k in enumerate(df_1.iloc[sample_idx]['cell_order'].split(' '))}\n    sample_2 = {k: v for v, k in enumerate(df_2.iloc[sample_idx]['cell_order'].split(' '))}\n    for key in sample_1: \n        sample_1[key] = ((sample_1[key] * 0.748) + (sample_2[key] * 0.252))\n    new_samples.append(' '.join([i[0] for i in list(sorted(sample_1.items(), key = lambda x: x[1]))]))\ndf_1['cell_order'] = new_samples","metadata":{"execution":{"iopub.status.busy":"2022-08-01T03:39:56.591672Z","iopub.execute_input":"2022-08-01T03:39:56.592229Z","iopub.status.idle":"2022-08-01T03:39:56.608547Z","shell.execute_reply.started":"2022-08-01T03:39:56.592179Z","shell.execute_reply":"2022-08-01T03:39:56.607220Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_1.to_csv('submission.csv', index = False)\ndf_1","metadata":{"execution":{"iopub.status.busy":"2022-08-01T03:39:56.610210Z","iopub.execute_input":"2022-08-01T03:39:56.610916Z","iopub.status.idle":"2022-08-01T03:39:56.622690Z","shell.execute_reply.started":"2022-08-01T03:39:56.610876Z","shell.execute_reply":"2022-08-01T03:39:56.621534Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}