{"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":"  def read_notebook(path):\n    import pandas as pd\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-15T17:16:09.613206Z","iopub.execute_input":"2022-07-15T17:16:09.613658Z","iopub.status.idle":"2022-07-15T17:16:09.648868Z","shell.execute_reply.started":"2022-07-15T17:16:09.613566Z","shell.execute_reply":"2022-07-15T17:16:09.647589Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def clean_code(cell): return str(cell).replace(\"\\\\n\", \"\\n\")\n\ndef sample_cells(cells, n):\n    import numpy as np\n    cells = [clean_code(cell) for cell in cells]\n    if n >= len(cells): return [cell[:200] for cell in cells]\n    else:\n        results = []\n        step = len(cells) / n\n        idx = 0\n        while int(np.round(idx)) < len(cells):\n            results.append(cells[int(np.round(idx))])\n            idx += step        \n        if cells[-1] not in results: results[-1] = cells[-1]\n        return results\ndef get_features(df):\n    from tqdm import tqdm\n    features = dict()\n    df = df.sort_values(\"rank\").reset_index(drop=True)\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 = 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    return features","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:09.698967Z","iopub.execute_input":"2022-07-15T17:16:09.699788Z","iopub.status.idle":"2022-07-15T17:16:09.713890Z","shell.execute_reply.started":"2022-07-15T17:16:09.699734Z","shell.execute_reply":"2022-07-15T17:16:09.712597Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def read_data(data): return tuple(d.cuda() for d in data[:-1]), data[-1].cuda()\n\ndef validate(model, val_loader):    \n    import sys\n    import torch    \n    import numpy as np\n    from tqdm import tqdm    \n    model.eval()    \n    tbar = tqdm(val_loader, file=sys.stdout)    \n    preds = []\n    labels = []\n    with torch.no_grad():\n        for idx, data in enumerate(tbar):\n            inputs, target = read_data(data)\n            pred = model(*inputs)\n            preds.append(pred.detach().cpu().numpy().ravel())\n            labels.append(target.detach().cpu().numpy().ravel())    \n    return np.concatenate(labels), np.concatenate(preds)\n\ndef predict_caller(args): return predict(args[0], args[1])\n    \ndef predict(model_path, ckpt_path):\n    import gc\n    import json\n    import sys, os\n    import numpy as np\n    import pandas as pd\n    from tqdm import tqdm\n    from pathlib import Path\n    from scipy import sparse\n\n    data_dir = Path('../input/AI4Code')\n    paths_test = list((data_dir / 'test').glob('*.json'))\n    notebooks_test = [\n        read_notebook(path) for path in tqdm(paths_test, desc='Test NBs')\n    ]\n    test_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    test_df[\"rank\"] = test_df.groupby([\"id\", \"cell_type\"]).cumcount()\n    test_df[\"pred\"] = test_df.groupby([\"id\", \"cell_type\"])[\"rank\"].rank(pct=True)\n    \n    import torch\n    import torch.nn as nn\n    import torch.nn.functional as F\n    from torch.utils.data import DataLoader, Dataset\n    from transformers import AutoModel, AutoTokenizer\n\n    class MarkdownModel(nn.Module):\n        def __init__(self, model_path):\n            super(MarkdownModel, self).__init__()\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 = self.top(torch.cat((x[:, 0, :], fts),1))\n            return x\n\n\n    class MarkdownDataset(Dataset):\n\n        def __init__(self, df, model_name_or_path, total_max_len, md_max_len, fts):\n            super().__init__()\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.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n            self.fts = fts\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            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_md = self.fts[row.id][\"total_md\"]\n            n_code = self.fts[row.id][\"total_md\"]\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) == self.total_max_len\n\n            return ids, mask, fts, torch.FloatTensor([row.pct_rank])\n            def __len__(self):\n                return self.df.shape[0]\n    \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 = 8\n    MAX_LEN = 64\n    test_df[\"pct_rank\"] = 0\n    test_fts = get_features(test_df)\n    test_ds = MarkdownDataset(test_df[test_df[\"cell_type\"] == \"markdown\"].reset_index(drop=True), md_max_len=64,total_max_len=512, model_name_or_path=model_path, 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    model.to(torch.device('cpu'))\n    torch.cuda.empty_cache()    \n    del model, test_loader, test_ds\n    gc.collect()      \n    \n    test_df.loc[test_df[\"cell_type\"] == \"markdown\", \"pred\"] = y_test\n    sub_df = test_df.sort_values(\"pred\").groupby(\"id\")[\"cell_id\"].apply(lambda x: \" \".join(x)).reset_index()\n    sub_df.rename(columns={\"cell_id\": \"cell_order\"}, inplace=True)\n    sub_df.head()\n    sub_df.to_csv(\"submission_1.csv\", index=False)\n\n    del test_df, paths_test, notebooks_test, test_fts, model_path, ckpt_path, sub_df\n    del json, np, pd, tqdm, Path, sparse, torch, sys, os, nn, F, AutoModel, AutoTokenizer\n    gc.collect()\n    ","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:09.825111Z","iopub.execute_input":"2022-07-15T17:16:09.825677Z","iopub.status.idle":"2022-07-15T17:16:09.870859Z","shell.execute_reply.started":"2022-07-15T17:16:09.825625Z","shell.execute_reply":"2022-07-15T17:16:09.869623Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import json\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom pathlib import Path\nfrom scipy import sparse\n\npd.options.display.width = 180\npd.options.display.max_colwidth = 120\n\nBERT_PATH = \"../input/huggingface-bert-variants/distilbert-base-uncased/distilbert-base-uncased\"\n\ndata_dir = Path('../input/AI4Code')\nNUM_TRAIN = 200\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         )\n\npaths_train = list((data_dir / 'train').glob('*.json'))[:NUM_TRAIN]\nnotebooks_train = [\n    read_notebook(path) for path in tqdm(paths_train, desc='Train NBs')\n]\ndf = (\n    pd.concat(notebooks_train)\n    .set_index('id', append=True)\n    .swaplevel()\n    .sort_index(level='id', sort_remaining=False)\n)\ndf","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:09.941434Z","iopub.execute_input":"2022-07-15T17:16:09.941938Z","iopub.status.idle":"2022-07-15T17:16:15.283165Z","shell.execute_reply.started":"2022-07-15T17:16:09.941888Z","shell.execute_reply":"2022-07-15T17:16:15.281804Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"nb_id = df.index.unique('id')[6]\nprint('Notebook:', nb_id)\n\nprint(\"The disordered notebook:\")\nnb = df.loc[nb_id, :]\ndisplay(nb)\nprint()","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:15.285282Z","iopub.execute_input":"2022-07-15T17:16:15.285931Z","iopub.status.idle":"2022-07-15T17:16:15.306859Z","shell.execute_reply.started":"2022-07-15T17:16:15.285860Z","shell.execute_reply":"2022-07-15T17:16:15.305725Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_orders = pd.read_csv(\n    data_dir / 'train_orders.csv',\n    index_col='id',\n    squeeze=True,\n).str.split()\n\ndf_orders","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:15.308391Z","iopub.execute_input":"2022-07-15T17:16:15.308844Z","iopub.status.idle":"2022-07-15T17:16:18.495407Z","shell.execute_reply.started":"2022-07-15T17:16:15.308796Z","shell.execute_reply":"2022-07-15T17:16:18.494435Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(df_orders.loc[\"002ba502bdac45\"])","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:18.497321Z","iopub.execute_input":"2022-07-15T17:16:18.498365Z","iopub.status.idle":"2022-07-15T17:16:18.506781Z","shell.execute_reply.started":"2022-07-15T17:16:18.498324Z","shell.execute_reply":"2022-07-15T17:16:18.505996Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cell_order = df_orders.loc[nb_id]\n\nprint(\"The ordered notebook:\")\nnb.loc[cell_order, :]","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:18.507869Z","iopub.execute_input":"2022-07-15T17:16:18.508859Z","iopub.status.idle":"2022-07-15T17:16:18.532718Z","shell.execute_reply.started":"2022-07-15T17:16:18.508823Z","shell.execute_reply":"2022-07-15T17:16:18.531290Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_ranks(base, derived):\n    return [base.index(d) for d in derived]\n\ncell_ranks = get_ranks(cell_order, list(nb.index))\nnb.insert(0, 'rank', cell_ranks)\n\nnb","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:18.534100Z","iopub.execute_input":"2022-07-15T17:16:18.534486Z","iopub.status.idle":"2022-07-15T17:16:18.551815Z","shell.execute_reply.started":"2022-07-15T17:16:18.534453Z","shell.execute_reply":"2022-07-15T17:16:18.550626Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_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_ranks","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:18.553057Z","iopub.execute_input":"2022-07-15T17:16:18.554007Z","iopub.status.idle":"2022-07-15T17:16:18.633497Z","shell.execute_reply.started":"2022-07-15T17:16:18.553966Z","shell.execute_reply":"2022-07-15T17:16:18.632379Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_ancestors = pd.read_csv(data_dir / 'train_ancestors.csv', index_col='id')\ndf_ancestors","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:18.634713Z","iopub.execute_input":"2022-07-15T17:16:18.635811Z","iopub.status.idle":"2022-07-15T17:16:18.894842Z","shell.execute_reply.started":"2022-07-15T17:16:18.635771Z","shell.execute_reply":"2022-07-15T17:16:18.893708Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = df.reset_index().merge(df_ranks, on=[\"id\", \"cell_id\"]).merge(df_ancestors, on=[\"id\"])\ndf","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:18.896412Z","iopub.execute_input":"2022-07-15T17:16:18.897098Z","iopub.status.idle":"2022-07-15T17:16:18.978103Z","shell.execute_reply.started":"2022-07-15T17:16:18.897063Z","shell.execute_reply":"2022-07-15T17:16:18.976767Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df[\"pct_rank\"] = df[\"rank\"] / df.groupby(\"id\")[\"cell_id\"].transform(\"count\")\ndf[\"pct_rank\"].hist(bins=8)","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:18.982898Z","iopub.execute_input":"2022-07-15T17:16:18.983233Z","iopub.status.idle":"2022-07-15T17:16:19.228008Z","shell.execute_reply.started":"2022-07-15T17:16:18.983205Z","shell.execute_reply":"2022-07-15T17:16:19.226675Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dict_cellid_source = dict(zip(df['cell_id'].values, df['source'].values))\nimport os\nimport re\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom pathlib import Path\nimport matplotlib.pyplot as plt\nfrom sklearn.metrics.pairwise import cosine_similarity\nfrom nltk.stem import WordNetLemmatizer\nimport nltk; nltk.download('wordnet')\n\nstemmer = WordNetLemmatizer()\n\ndef preprocess_text(document):\n        document = re.sub(r'\\W', ' ', str(document))\n\n        document = re.sub(r'\\s+[a-zA-Z]\\s+', ' ', document)\n\n        document = re.sub(r'\\^[a-zA-Z]\\s+', ' ', document)\n\n        document = re.sub(r'\\s+', ' ', document, flags=re.I)\n\n        document = re.sub(r'^b\\s+', '', document)\n\n        document = document.lower()\n\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\n        preprocessed_text = ' '.join(tokens)\n        return preprocessed_text\n\n    \ndef preprocess_df(df):\n    \"\"\"\n    This function is for processing sorce of notebook\n    returns preprocessed dataframe\n    \"\"\"\n    return [preprocess_text(message) for message in df.source]\n\ndf.source = df.source.apply(preprocess_text)","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:19.230021Z","iopub.execute_input":"2022-07-15T17:16:19.230456Z","iopub.status.idle":"2022-07-15T17:16:23.933747Z","shell.execute_reply.started":"2022-07-15T17:16:19.230409Z","shell.execute_reply":"2022-07-15T17:16:23.932553Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.model_selection import GroupShuffleSplit\n\nNVALID = 0.1  \n\nsplitter = GroupShuffleSplit(n_splits=1, test_size=NVALID, random_state=0)\n\ntrain_ind, val_ind = next(splitter.split(df, groups=df[\"ancestor_id\"]))\n\ntrain_df = df.loc[train_ind].reset_index(drop=True)\nval_df = df.loc[val_ind].reset_index(drop=True)","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:23.934939Z","iopub.execute_input":"2022-07-15T17:16:23.935235Z","iopub.status.idle":"2022-07-15T17:16:23.955023Z","shell.execute_reply.started":"2022-07-15T17:16:23.935207Z","shell.execute_reply":"2022-07-15T17:16:23.953804Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tqdm.notebook import tqdm\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          \n        elif mode == 'test':\n          triplets.append( [cell_id, cid, label] )\n        elif random_drop[count%10000]:\n          triplets.append( [cell_id, cid, label] )\n\n    \n  return triplets\n\ntriplets = generate_triplet(train_df)\nval_triplets = generate_triplet(val_df, mode = 'test')","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:23.956475Z","iopub.execute_input":"2022-07-15T17:16:23.956946Z","iopub.status.idle":"2022-07-15T17:16:24.858410Z","shell.execute_reply.started":"2022-07-15T17:16:23.956899Z","shell.execute_reply":"2022-07-15T17:16:24.857311Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"val_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:24.860447Z","iopub.execute_input":"2022-07-15T17:16:24.861085Z","iopub.status.idle":"2022-07-15T17:16:24.876848Z","shell.execute_reply.started":"2022-07-15T17:16:24.861051Z","shell.execute_reply":"2022-07-15T17:16:24.875302Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from bisect import bisect\n\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 \n    for gt, pred in zip(ground_truth, predictions):\n        ranks = [gt.index(x) for x in pred] \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-15T17:16:24.878128Z","iopub.execute_input":"2022-07-15T17:16:24.878510Z","iopub.status.idle":"2022-07-15T17:16:24.893304Z","shell.execute_reply.started":"2022-07-15T17:16:24.878475Z","shell.execute_reply":"2022-07-15T17:16:24.892070Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from transformers import RobertaTokenizer, RobertaConfig, RobertaModel\nimport torch.nn.functional as F\nimport torch.nn as nn\nimport torch\nfrom transformers import AutoModelWithLMHead, AutoTokenizer, AutoModel\n\nMAX_LEN = 128\n\n    \nclass MarkdownModel(nn.Module):\n    def __init__(self):\n        super(MarkdownModel, self).__init__()\n        self.distill_bert = AutoModel.from_pretrained(\"../input/mymodelpairbertsmallpretrained/models/checkpoint-18000\")\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","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:24.895285Z","iopub.execute_input":"2022-07-15T17:16:24.896400Z","iopub.status.idle":"2022-07-15T17:16:27.606884Z","shell.execute_reply.started":"2022-07-15T17:16:24.896352Z","shell.execute_reply":"2022-07-15T17:16:27.605950Z"},"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 = 1e-3\n    elif epoch < 5:\n        lr = 1e-4\n    else:\n        lr = 1e-5\n\n    for p in optimizer.param_groups:\n        p['lr'] = lr\n    return lr\n    \ndef get_optimizer(net):\n    optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=3e-4, betas=(0.9, 0.999), eps=1e-08)\n    return optimizer\n\nBS = 128\nNW = 8","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:27.608294Z","iopub.execute_input":"2022-07-15T17:16:27.609089Z","iopub.status.idle":"2022-07-15T17:16:27.616380Z","shell.execute_reply.started":"2022-07-15T17:16:27.609054Z","shell.execute_reply":"2022-07-15T17:16:27.615093Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def read_data(data):\n    return tuple(d.cuda() for d in data[:-1]), data[-1].cuda()\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)","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:27.619382Z","iopub.execute_input":"2022-07-15T17:16:27.620393Z","iopub.status.idle":"2022-07-15T17:16:27.635508Z","shell.execute_reply.started":"2022-07-15T17:16:27.620340Z","shell.execute_reply":"2022-07-15T17:16:27.634310Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"paths_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()","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:27.637258Z","iopub.execute_input":"2022-07-15T17:16:27.638645Z","iopub.status.idle":"2022-07-15T17:16:27.741230Z","shell.execute_reply.started":"2022-07-15T17:16:27.638594Z","shell.execute_reply":"2022-07-15T17:16:27.740035Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_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')","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:27.742716Z","iopub.execute_input":"2022-07-15T17:16:27.743114Z","iopub.status.idle":"2022-07-15T17:16:27.847574Z","shell.execute_reply.started":"2022-07-15T17:16:27.743081Z","shell.execute_reply":"2022-07-15T17:16:27.846595Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_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\n  count += N_mark * N_code","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:27.848909Z","iopub.execute_input":"2022-07-15T17:16:27.849430Z","iopub.status.idle":"2022-07-15T17:16:27.908033Z","shell.execute_reply.started":"2022-07-15T17:16:27.849397Z","shell.execute_reply":"2022-07-15T17:16:27.906738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub_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.to_csv(\"submission_2.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:27.909672Z","iopub.execute_input":"2022-07-15T17:16:27.909982Z","iopub.status.idle":"2022-07-15T17:16:27.922506Z","shell.execute_reply.started":"2022-07-15T17:16:27.909954Z","shell.execute_reply":"2022-07-15T17:16:27.921691Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.read_csv('submission_2.csv')","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:27.924717Z","iopub.execute_input":"2022-07-15T17:16:27.925508Z","iopub.status.idle":"2022-07-15T17:16:27.934965Z","shell.execute_reply.started":"2022-07-15T17:16:27.925459Z","shell.execute_reply":"2022-07-15T17:16:27.933715Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"new_samples = []\nfor sample_idx in range(len(df)):\n    sample_1 = {k: v for v, k in enumerate(df.iloc[sample_idx]['cell_order'].split(' '))}\n    for key in sample_1: sample_1[key] = ( (sample_1[key] * 0.4))\n    new_samples.append(' '.join([i[0] for i in list(sorted(sample_1.items(), key=lambda x:x[1]))]))\ndf['cell_order'] = new_samples","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:27.936888Z","iopub.execute_input":"2022-07-15T17:16:27.937746Z","iopub.status.idle":"2022-07-15T17:16:27.949263Z","shell.execute_reply.started":"2022-07-15T17:16:27.937710Z","shell.execute_reply":"2022-07-15T17:16:27.948072Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df.to_csv('submission.csv', index = False)\ndf","metadata":{"execution":{"iopub.status.busy":"2022-07-15T17:16:27.952103Z","iopub.execute_input":"2022-07-15T17:16:27.952631Z","iopub.status.idle":"2022-07-15T17:16:27.976154Z","shell.execute_reply.started":"2022-07-15T17:16:27.952585Z","shell.execute_reply":"2022-07-15T17:16:27.974909Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import json\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nfrom scipy import sparse\nfrom tqdm import tqdm\nfrom sklearn.feature_extraction.text import TfidfVectorizer\n\npd.options.display.width = 180\npd.options.display.max_colwidth = 120","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:18:58.863670Z","iopub.execute_input":"2022-07-23T08:18:58.864183Z","iopub.status.idle":"2022-07-23T08:18:59.384334Z","shell.execute_reply.started":"2022-07-23T08:18:58.864083Z","shell.execute_reply":"2022-07-23T08:18:59.383295Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_dir = Path('../input/AI4Code')","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:19:10.850703Z","iopub.execute_input":"2022-07-23T08:19:10.851115Z","iopub.status.idle":"2022-07-23T08:19:10.857051Z","shell.execute_reply.started":"2022-07-23T08:19:10.851081Z","shell.execute_reply":"2022-07-23T08:19:10.855951Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"NUM_TRAIN = 10000\n\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    )\n\n\npaths_train = list((data_dir / 'train').glob('*.json'))[:NUM_TRAIN]\nnotebooks_train = [\n    read_notebook(path) for path in tqdm(paths_train, desc='Train NBs')\n]\ndf = (\n    pd.concat(notebooks_train)\n    .set_index('id', append=True)\n    .swaplevel()\n    .sort_index(level='id', sort_remaining=False)\n)\ndf","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:19:37.069959Z","iopub.execute_input":"2022-07-23T08:19:37.070370Z","iopub.status.idle":"2022-07-23T08:21:05.420048Z","shell.execute_reply.started":"2022-07-23T08:19:37.070334Z","shell.execute_reply":"2022-07-23T08:21:05.418874Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import plotly.io as pio\npio.renderers.default='notebook'\nimport plotly.express as px\n\ndf_temp = df.reset_index()\npie_data = df_temp[\"cell_type\"].value_counts().reset_index()\npie_data.columns = [\"cell_type\", \"count\"]\n\nfig = px.pie(pie_data, values='count', names='cell_type', title='Code vs Markdown')\nfig.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:21:23.903061Z","iopub.execute_input":"2022-07-23T08:21:23.903439Z","iopub.status.idle":"2022-07-23T08:21:26.025440Z","shell.execute_reply.started":"2022-07-23T08:21:23.903409Z","shell.execute_reply":"2022-07-23T08:21:26.024409Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cell_analysis = df_temp.groupby([\"id\", \"cell_type\"])[\"cell_id\"].count().reset_index()\nscatter_data = pd.pivot(data=cell_analysis, index=\"id\", columns=\"cell_type\", values=\"cell_id\")\nscatter_data[\"size\"] = 30\n\nfig = px.scatter(scatter_data, x=\"code\", y=\"markdown\")\nfig.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:21:46.427206Z","iopub.execute_input":"2022-07-23T08:21:46.428432Z","iopub.status.idle":"2022-07-23T08:21:46.737511Z","shell.execute_reply.started":"2022-07-23T08:21:46.428378Z","shell.execute_reply":"2022-07-23T08:21:46.736423Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"nb_id = df.index.unique('id')[6]\nprint('Notebook:', nb_id)\n\nprint(\"The disordered notebook:\")\nnb = df.loc[nb_id, :]\ndisplay(nb)\nprint()","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:22:07.134274Z","iopub.execute_input":"2022-07-23T08:22:07.134687Z","iopub.status.idle":"2022-07-23T08:22:07.161464Z","shell.execute_reply.started":"2022-07-23T08:22:07.134653Z","shell.execute_reply":"2022-07-23T08:22:07.160744Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_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.head(10)","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:22:25.843363Z","iopub.execute_input":"2022-07-23T08:22:25.843751Z","iopub.status.idle":"2022-07-23T08:22:28.990632Z","shell.execute_reply.started":"2022-07-23T08:22:25.843698Z","shell.execute_reply":"2022-07-23T08:22:28.989481Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cell_order = df_orders.loc[nb_id]\n\nprint(\"The ordered notebook:\")\nnb.loc[cell_order, :]","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:22:49.924145Z","iopub.execute_input":"2022-07-23T08:22:49.924893Z","iopub.status.idle":"2022-07-23T08:22:49.961849Z","shell.execute_reply.started":"2022-07-23T08:22:49.924853Z","shell.execute_reply":"2022-07-23T08:22:49.960787Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_ranks(base, derived):\n    return [base.index(d) for d in derived]\n\ncell_ranks = get_ranks(cell_order, list(nb.index))\nnb.insert(0, 'rank', cell_ranks)\n\nnb","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:23:16.566678Z","iopub.execute_input":"2022-07-23T08:23:16.567100Z","iopub.status.idle":"2022-07-23T08:23:16.583364Z","shell.execute_reply.started":"2022-07-23T08:23:16.567064Z","shell.execute_reply":"2022-07-23T08:23:16.582479Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pandas.testing import assert_frame_equal\n\nassert_frame_equal(nb.loc[cell_order, :], nb.sort_values('rank'))","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:23:37.012793Z","iopub.execute_input":"2022-07-23T08:23:37.013186Z","iopub.status.idle":"2022-07-23T08:23:37.023216Z","shell.execute_reply.started":"2022-07-23T08:23:37.013154Z","shell.execute_reply":"2022-07-23T08:23:37.022334Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_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_ranks","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:23:53.184094Z","iopub.execute_input":"2022-07-23T08:23:53.184515Z","iopub.status.idle":"2022-07-23T08:23:55.716945Z","shell.execute_reply.started":"2022-07-23T08:23:53.184480Z","shell.execute_reply":"2022-07-23T08:23:55.715881Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_ancestors = pd.read_csv(data_dir / 'train_ancestors.csv', index_col='id')\ndf_ancestors","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:24:10.774035Z","iopub.execute_input":"2022-07-23T08:24:10.774449Z","iopub.status.idle":"2022-07-23T08:24:11.028114Z","shell.execute_reply.started":"2022-07-23T08:24:10.774416Z","shell.execute_reply":"2022-07-23T08:24:11.026942Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.model_selection import GroupShuffleSplit\n\nNVALID = 0.1  \n\nsplitter = GroupShuffleSplit(n_splits=1, test_size=NVALID, random_state=0)\n\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, :]","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:24:41.516968Z","iopub.execute_input":"2022-07-23T08:24:41.517770Z","iopub.status.idle":"2022-07-23T08:24:41.694844Z","shell.execute_reply.started":"2022-07-23T08:24:41.517731Z","shell.execute_reply":"2022-07-23T08:24:41.693048Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.feature_extraction.text import TfidfVectorizer\n\ntfidf = TfidfVectorizer(min_df=0.01)\nX_train = tfidf.fit_transform(df_train['source'].astype(str))\ny_train = df_ranks.loc[ids_train].to_numpy()\ngroups = df_ranks.loc[ids_train].groupby('id').size().to_numpy()","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:25:15.496847Z","iopub.execute_input":"2022-07-23T08:25:15.497798Z","iopub.status.idle":"2022-07-23T08:25:28.803421Z","shell.execute_reply.started":"2022-07-23T08:25:15.497756Z","shell.execute_reply":"2022-07-23T08:25:28.802625Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X_train = sparse.hstack((\n    X_train,\n    np.where(\n        df_train['cell_type'] == 'code',\n        df_train.groupby(['id', 'cell_type']).cumcount().to_numpy() + 1,\n        0,\n    ).reshape(-1, 1)\n))\nprint(X_train.shape)","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:25:36.234135Z","iopub.execute_input":"2022-07-23T08:25:36.234531Z","iopub.status.idle":"2022-07-23T08:25:36.493006Z","shell.execute_reply.started":"2022-07-23T08:25:36.234498Z","shell.execute_reply":"2022-07-23T08:25:36.491743Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from xgboost import XGBRanker\n\nmodel = XGBRanker(\n    min_child_weight=10,\n    subsample=0.5,\n    tree_method='hist',\n)\nmodel.fit(X_train, y_train, group=groups)","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:25:53.377295Z","iopub.execute_input":"2022-07-23T08:25:53.378339Z","iopub.status.idle":"2022-07-23T08:26:05.888689Z","shell.execute_reply.started":"2022-07-23T08:25:53.378283Z","shell.execute_reply":"2022-07-23T08:26:05.887528Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X_valid = tfidf.transform(df_valid['source'].astype(str))\n\ny_valid = df_orders.loc[ids_valid]\n\nX_valid = sparse.hstack((\n    X_valid,\n    np.where(\n        df_valid['cell_type'] == 'code',\n        df_valid.groupby(['id', 'cell_type']).cumcount().to_numpy() + 1,\n        0,\n    ).reshape(-1, 1)\n))","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:26:16.839267Z","iopub.execute_input":"2022-07-23T08:26:16.840199Z","iopub.status.idle":"2022-07-23T08:26:18.168367Z","shell.execute_reply.started":"2022-07-23T08:26:16.840158Z","shell.execute_reply":"2022-07-23T08:26:18.167226Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from catboost import FeaturesData","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:26:35.252817Z","iopub.execute_input":"2022-07-23T08:26:35.253181Z","iopub.status.idle":"2022-07-23T08:26:35.550915Z","shell.execute_reply.started":"2022-07-23T08:26:35.253152Z","shell.execute_reply":"2022-07-23T08:26:35.549663Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"nb_id = df.index.unique('id')[6]\nprint('Notebook:', nb_id)\n\nprint(\"The disordered notebook:\")\nnb = df.loc[nb_id, :]\ndisplay(nb)\nprint()","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:26:52.715860Z","iopub.execute_input":"2022-07-23T08:26:52.716274Z","iopub.status.idle":"2022-07-23T08:26:52.738470Z","shell.execute_reply.started":"2022-07-23T08:26:52.716240Z","shell.execute_reply":"2022-07-23T08:26:52.737669Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_orders = pd.read_csv(\n    data_dir / 'train_orders.csv',\n    index_col='id',\n    squeeze=True,\n).str.split() \n\ndf_orders","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:27:18.024951Z","iopub.execute_input":"2022-07-23T08:27:18.025925Z","iopub.status.idle":"2022-07-23T08:27:20.049155Z","shell.execute_reply.started":"2022-07-23T08:27:18.025876Z","shell.execute_reply":"2022-07-23T08:27:20.047969Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cell_order = df_orders.loc[nb_id]\n\nprint(\"The ordered notebook:\")\nnb.loc[cell_order, :]","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:27:39.279822Z","iopub.execute_input":"2022-07-23T08:27:39.280231Z","iopub.status.idle":"2022-07-23T08:27:39.297740Z","shell.execute_reply.started":"2022-07-23T08:27:39.280196Z","shell.execute_reply":"2022-07-23T08:27:39.296886Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_ranks(base, derived):\n    return [base.index(d) for d in derived]\n\ncell_ranks = get_ranks(cell_order, list(nb.index))\nnb.insert(0, 'rank', cell_ranks)\n\nnb","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:28:02.155749Z","iopub.execute_input":"2022-07-23T08:28:02.156150Z","iopub.status.idle":"2022-07-23T08:28:02.172541Z","shell.execute_reply.started":"2022-07-23T08:28:02.156118Z","shell.execute_reply":"2022-07-23T08:28:02.171779Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pandas.testing import assert_frame_equal\n\nassert_frame_equal(nb.loc[cell_order, :], nb.sort_values('rank'))","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:28:19.097150Z","iopub.execute_input":"2022-07-23T08:28:19.097565Z","iopub.status.idle":"2022-07-23T08:28:19.106797Z","shell.execute_reply.started":"2022-07-23T08:28:19.097530Z","shell.execute_reply":"2022-07-23T08:28:19.105623Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_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_ranks","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:28:37.365934Z","iopub.execute_input":"2022-07-23T08:28:37.366337Z","iopub.status.idle":"2022-07-23T08:28:39.649048Z","shell.execute_reply.started":"2022-07-23T08:28:37.366303Z","shell.execute_reply":"2022-07-23T08:28:39.648018Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_ancestors = pd.read_csv(data_dir / 'train_ancestors.csv', index_col='id')\ndf_ancestors","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:28:58.005256Z","iopub.execute_input":"2022-07-23T08:28:58.005849Z","iopub.status.idle":"2022-07-23T08:28:58.194411Z","shell.execute_reply.started":"2022-07-23T08:28:58.005798Z","shell.execute_reply":"2022-07-23T08:28:58.193368Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.model_selection import GroupShuffleSplit\n\nNVALID = 0.1  \nsplitter = GroupShuffleSplit(n_splits=1, test_size=NVALID, random_state=0)\n\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, :]","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:29:26.030408Z","iopub.execute_input":"2022-07-23T08:29:26.030837Z","iopub.status.idle":"2022-07-23T08:29:26.156138Z","shell.execute_reply.started":"2022-07-23T08:29:26.030799Z","shell.execute_reply":"2022-07-23T08:29:26.155182Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.feature_extraction.text import TfidfVectorizer\n\ntfidf = TfidfVectorizer(min_df=0.01)\nX_train = tfidf.fit_transform(df_train['source'].astype(str))\ny_train = df_ranks.loc[ids_train].to_numpy()\ngroups = df_ranks.loc[ids_train].groupby('id').size().to_numpy()","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:29:56.822550Z","iopub.execute_input":"2022-07-23T08:29:56.823322Z","iopub.status.idle":"2022-07-23T08:30:09.882184Z","shell.execute_reply.started":"2022-07-23T08:29:56.823280Z","shell.execute_reply":"2022-07-23T08:30:09.881248Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X_train = sparse.hstack((\n    X_train,\n    np.where(\n        df_train['cell_type'] == 'code',\n        df_train.groupby(['id', 'cell_type']).cumcount().to_numpy() + 1,\n        0,\n    ).reshape(-1, 1)\n))\nprint(X_train.shape)","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:30:20.359061Z","iopub.execute_input":"2022-07-23T08:30:20.360039Z","iopub.status.idle":"2022-07-23T08:30:20.596860Z","shell.execute_reply.started":"2022-07-23T08:30:20.359988Z","shell.execute_reply":"2022-07-23T08:30:20.595646Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from xgboost import XGBRanker\n\nmodel = XGBRanker(\n    min_child_weight=10,\n    subsample=0.5,\n    tree_method='hist',\n)\nmodel.fit(X_train, y_train, group=groups)","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:30:40.035280Z","iopub.execute_input":"2022-07-23T08:30:40.035798Z","iopub.status.idle":"2022-07-23T08:30:52.183604Z","shell.execute_reply.started":"2022-07-23T08:30:40.035752Z","shell.execute_reply":"2022-07-23T08:30:52.182579Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X_valid = tfidf.transform(df_valid['source'].astype(str))\ny_valid = df_orders.loc[ids_valid]\n\nX_valid = sparse.hstack((\n    X_valid,\n    np.where(\n        df_valid['cell_type'] == 'code',\n        df_valid.groupby(['id', 'cell_type']).cumcount().to_numpy() + 1,\n        0,\n    ).reshape(-1, 1)\n))","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:31:00.184097Z","iopub.execute_input":"2022-07-23T08:31:00.184484Z","iopub.status.idle":"2022-07-23T08:31:01.540695Z","shell.execute_reply.started":"2022-07-23T08:31:00.184452Z","shell.execute_reply":"2022-07-23T08:31:01.539767Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"y_pred = pd.DataFrame({'rank': model.predict(X_valid)}, index=df_valid.index)\ny_pred = (\n    y_pred\n    .sort_values(['id', 'rank'])\n    .reset_index('cell_id')  \n    .groupby('id')['cell_id'].apply(list) \n)\ny_pred.head(10)","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:31:43.715273Z","iopub.execute_input":"2022-07-23T08:31:43.715660Z","iopub.status.idle":"2022-07-23T08:31:43.895750Z","shell.execute_reply.started":"2022-07-23T08:31:43.715628Z","shell.execute_reply":"2022-07-23T08:31:43.894764Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"nb_id = df_valid.index.get_level_values('id').unique()[8]\n\ndisplay(df.loc[nb_id])\ndisplay(df.loc[nb_id].loc[y_pred.loc[nb_id]])","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:32:03.208209Z","iopub.execute_input":"2022-07-23T08:32:03.208621Z","iopub.status.idle":"2022-07-23T08:32:03.238191Z","shell.execute_reply.started":"2022-07-23T08:32:03.208578Z","shell.execute_reply":"2022-07-23T08:32:03.237288Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from bisect import bisect\n\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 \n    for gt, pred in zip(ground_truth, predictions):\n        ranks = [gt.index(x) for x in pred] \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-23T08:32:48.181493Z","iopub.execute_input":"2022-07-23T08:32:48.181939Z","iopub.status.idle":"2022-07-23T08:32:48.190351Z","shell.execute_reply.started":"2022-07-23T08:32:48.181893Z","shell.execute_reply":"2022-07-23T08:32:48.189165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"y_dummy = df_valid.reset_index('cell_id').groupby('id')['cell_id'].apply(list)\nkendall_tau(y_valid, y_dummy)","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:33:13.310866Z","iopub.execute_input":"2022-07-23T08:33:13.311227Z","iopub.status.idle":"2022-07-23T08:33:13.427394Z","shell.execute_reply.started":"2022-07-23T08:33:13.311199Z","shell.execute_reply":"2022-07-23T08:33:13.425671Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"kendall_tau(y_valid, y_pred)","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:33:26.212766Z","iopub.execute_input":"2022-07-23T08:33:26.213823Z","iopub.status.idle":"2022-07-23T08:33:26.286787Z","shell.execute_reply.started":"2022-07-23T08:33:26.213786Z","shell.execute_reply":"2022-07-23T08:33:26.285977Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"paths_test = list((data_dir / 'test').glob('*.json'))\nnotebooks_test = [\n    read_notebook(path) for path in tqdm(paths_test, desc='Test NBs')\n]\ndf_test = (\n    pd.concat(notebooks_test)\n    .set_index('id', append=True)\n    .swaplevel()\n    .sort_index(level='id', sort_remaining=False)\n)","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:33:47.383164Z","iopub.execute_input":"2022-07-23T08:33:47.383543Z","iopub.status.idle":"2022-07-23T08:33:47.431994Z","shell.execute_reply.started":"2022-07-23T08:33:47.383512Z","shell.execute_reply":"2022-07-23T08:33:47.430803Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X_test = tfidf.transform(df_test['source'].astype(str))\nX_test = sparse.hstack((\n    X_test,\n    np.where(\n        df_test['cell_type'] == 'code',\n        df_test.groupby(['id', 'cell_type']).cumcount().to_numpy() + 1,\n        0,\n    ).reshape(-1, 1)\n))","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:34:05.807173Z","iopub.execute_input":"2022-07-23T08:34:05.807578Z","iopub.status.idle":"2022-07-23T08:34:05.825162Z","shell.execute_reply.started":"2022-07-23T08:34:05.807542Z","shell.execute_reply":"2022-07-23T08:34:05.824094Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"y_infer = pd.DataFrame({'rank': model.predict(X_test)}, index=df_test.index)\ny_infer = y_infer.sort_values(['id', 'rank']).reset_index('cell_id').groupby('id')['cell_id'].apply(list)\ny_infer","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:34:16.700987Z","iopub.execute_input":"2022-07-23T08:34:16.701389Z","iopub.status.idle":"2022-07-23T08:34:16.725067Z","shell.execute_reply.started":"2022-07-23T08:34:16.701355Z","shell.execute_reply":"2022-07-23T08:34:16.724106Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"y_sample = pd.read_csv(data_dir / 'sample_submission.csv', index_col='id', squeeze=True)\ny_sample","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:34:38.942005Z","iopub.execute_input":"2022-07-23T08:34:38.942389Z","iopub.status.idle":"2022-07-23T08:34:38.956394Z","shell.execute_reply.started":"2022-07-23T08:34:38.942358Z","shell.execute_reply":"2022-07-23T08:34:38.955596Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"y_submit = (\n    y_infer\n    .apply(' '.join)  \n    .rename_axis('id')\n    .rename('cell_order')\n)\ny_submit","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:35:09.475641Z","iopub.execute_input":"2022-07-23T08:35:09.476090Z","iopub.status.idle":"2022-07-23T08:35:09.484232Z","shell.execute_reply.started":"2022-07-23T08:35:09.476055Z","shell.execute_reply":"2022-07-23T08:35:09.483378Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"y_submit.to_csv('submission.csv')","metadata":{"execution":{"iopub.status.busy":"2022-07-23T08:35:24.443974Z","iopub.execute_input":"2022-07-23T08:35:24.445010Z","iopub.status.idle":"2022-07-23T08:35:24.453039Z","shell.execute_reply.started":"2022-07-23T08:35:24.444960Z","shell.execute_reply":"2022-07-23T08:35:24.451963Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}