{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":35887,"databundleVersionId":3659326,"sourceType":"competition"},{"sourceId":4250,"sourceType":"modelInstanceVersion","modelInstanceId":3046,"modelId":575}],"dockerImageVersionId":30185,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# AI4Code Pytorch T5 Imlimentation\n\nBased on solution: https://www.kaggle.com/code/aerdem4/ai4code-pytorch-distilbert-baseline\nImplemented T5 model\n","metadata":{"papermill":{"duration":0.031568,"end_time":"2022-05-12T10:15:13.890382","exception":false,"start_time":"2022-05-12T10:15:13.858814","status":"completed"},"tags":[]}},{"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\n\npd.options.display.width = 180\npd.options.display.max_colwidth = 120\n\n#BERT_PATH = \"../input/huggingface-bert-variants/distilbert-base-uncased/distilbert-base-uncased\"\nT5_PATH = \"/kaggle/input/flan-t5/pytorch/small/2\"\n\ndata_dir = Path('../input/AI4Code')","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:55:23.967251Z","iopub.execute_input":"2025-05-11T11:55:23.968022Z","iopub.status.idle":"2025-05-11T11:55:24.010523Z","shell.execute_reply.started":"2025-05-11T11:55:23.967933Z","shell.execute_reply":"2025-05-11T11:55:24.009986Z"},"papermill":{"duration":0.122804,"end_time":"2022-05-12T10:15:14.04297","exception":false,"start_time":"2022-05-12T10:15:13.920166","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"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)\n\ndf","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:55:24.011705Z","iopub.execute_input":"2025-05-11T11:55:24.011914Z","iopub.status.idle":"2025-05-11T11:56:51.803163Z","shell.execute_reply.started":"2025-05-11T11:55:24.011890Z","shell.execute_reply":"2025-05-11T11:56:51.802496Z"},"papermill":{"duration":82.291505,"end_time":"2022-05-12T10:16:36.365197","exception":false,"start_time":"2022-05-12T10:15:14.073692","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Get an example notebook\nnb_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":"2025-05-11T11:56:51.804039Z","iopub.execute_input":"2025-05-11T11:56:51.804200Z","iopub.status.idle":"2025-05-11T11:56:51.823767Z","shell.execute_reply.started":"2025-05-11T11:56:51.804179Z","shell.execute_reply":"2025-05-11T11:56:51.823074Z"},"papermill":{"duration":0.270693,"end_time":"2022-05-12T10:16:36.882443","exception":false,"start_time":"2022-05-12T10:16:36.61175","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"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","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:56:51.825595Z","iopub.execute_input":"2025-05-11T11:56:51.826204Z","iopub.status.idle":"2025-05-11T11:56:54.197352Z","shell.execute_reply.started":"2025-05-11T11:56:51.826166Z","shell.execute_reply":"2025-05-11T11:56:54.196714Z"},"papermill":{"duration":2.835076,"end_time":"2022-05-12T10:16:39.9675","exception":false,"start_time":"2022-05-12T10:16:37.132424","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"len(df_orders.loc[\"002ba502bdac45\"])","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:56:54.198514Z","iopub.execute_input":"2025-05-11T11:56:54.198783Z","iopub.status.idle":"2025-05-11T11:56:54.206332Z","shell.execute_reply.started":"2025-05-11T11:56:54.198747Z","shell.execute_reply":"2025-05-11T11:56:54.205624Z"},"papermill":{"duration":0.257536,"end_time":"2022-05-12T10:16:40.472139","exception":false,"start_time":"2022-05-12T10:16:40.214603","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"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":"2025-05-11T11:56:54.207255Z","iopub.execute_input":"2025-05-11T11:56:54.207448Z","iopub.status.idle":"2025-05-11T11:56:54.224685Z","shell.execute_reply.started":"2025-05-11T11:56:54.207426Z","shell.execute_reply":"2025-05-11T11:56:54.224030Z"},"papermill":{"duration":0.265934,"end_time":"2022-05-12T10:16:40.98571","exception":false,"start_time":"2022-05-12T10:16:40.719776","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"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":"2025-05-11T11:56:54.225647Z","iopub.execute_input":"2025-05-11T11:56:54.225881Z","iopub.status.idle":"2025-05-11T11:56:54.240173Z","shell.execute_reply.started":"2025-05-11T11:56:54.225842Z","shell.execute_reply":"2025-05-11T11:56:54.239466Z"},"papermill":{"duration":0.265625,"end_time":"2022-05-12T10:16:41.501618","exception":false,"start_time":"2022-05-12T10:16:41.235993","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"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":"2025-05-11T11:56:54.241052Z","iopub.execute_input":"2025-05-11T11:56:54.241281Z","iopub.status.idle":"2025-05-11T11:56:56.186294Z","shell.execute_reply.started":"2025-05-11T11:56:54.241247Z","shell.execute_reply":"2025-05-11T11:56:56.185476Z"},"papermill":{"duration":2.967892,"end_time":"2022-05-12T10:16:44.752979","exception":false,"start_time":"2022-05-12T10:16:41.785087","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_ancestors = pd.read_csv(data_dir / 'train_ancestors.csv', index_col='id')\ndf_ancestors","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:56:56.187278Z","iopub.execute_input":"2025-05-11T11:56:56.187480Z","iopub.status.idle":"2025-05-11T11:56:56.355216Z","shell.execute_reply.started":"2025-05-11T11:56:56.187455Z","shell.execute_reply":"2025-05-11T11:56:56.354624Z"},"papermill":{"duration":0.44203,"end_time":"2022-05-12T10:16:45.446006","exception":false,"start_time":"2022-05-12T10:16:45.003976","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"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":"2025-05-11T11:56:56.358146Z","iopub.execute_input":"2025-05-11T11:56:56.358327Z","iopub.status.idle":"2025-05-11T11:56:57.014763Z","shell.execute_reply.started":"2025-05-11T11:56:56.358306Z","shell.execute_reply":"2025-05-11T11:56:57.014052Z"},"papermill":{"duration":1.007951,"end_time":"2022-05-12T10:16:46.70626","exception":false,"start_time":"2022-05-12T10:16:45.698309","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df[\"pct_rank\"] = df[\"rank\"] / df.groupby(\"id\")[\"cell_id\"].transform(\"count\")\n\ndf[\"pct_rank\"].hist(bins=10)","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:56:57.015623Z","iopub.execute_input":"2025-05-11T11:56:57.015853Z","iopub.status.idle":"2025-05-11T11:56:57.491801Z","shell.execute_reply.started":"2025-05-11T11:56:57.015828Z","shell.execute_reply":"2025-05-11T11:56:57.491076Z"},"papermill":{"duration":0.862186,"end_time":"2022-05-12T10:16:47.820945","exception":false,"start_time":"2022-05-12T10:16:46.958759","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.model_selection import GroupShuffleSplit\n\nNVALID = 0.1  # size of validation set\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":"2025-05-11T11:56:57.492652Z","iopub.execute_input":"2025-05-11T11:56:57.492871Z","iopub.status.idle":"2025-05-11T11:56:58.545377Z","shell.execute_reply.started":"2025-05-11T11:56:57.492846Z","shell.execute_reply":"2025-05-11T11:56:58.544798Z"},"papermill":{"duration":1.895199,"end_time":"2022-05-12T10:16:49.969199","exception":false,"start_time":"2022-05-12T10:16:48.074","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"val_df.head()","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:56:58.546339Z","iopub.execute_input":"2025-05-11T11:56:58.546555Z","iopub.status.idle":"2025-05-11T11:56:58.557384Z","shell.execute_reply.started":"2025-05-11T11:56:58.546530Z","shell.execute_reply":"2025-05-11T11:56:58.556726Z"},"papermill":{"duration":0.27219,"end_time":"2022-05-12T10:16:50.495502","exception":false,"start_time":"2022-05-12T10:16:50.223312","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"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  # 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":"2025-05-11T11:56:58.558207Z","iopub.execute_input":"2025-05-11T11:56:58.558415Z","iopub.status.idle":"2025-05-11T11:56:58.567824Z","shell.execute_reply.started":"2025-05-11T11:56:58.558380Z","shell.execute_reply":"2025-05-11T11:56:58.567194Z"},"papermill":{"duration":0.262837,"end_time":"2022-05-12T10:16:51.011588","exception":false,"start_time":"2022-05-12T10:16:50.748751","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"y_dummy = val_df.groupby('id')['cell_id'].apply(list)\nkendall_tau(df_orders.loc[y_dummy.index], y_dummy)","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:56:58.568655Z","iopub.execute_input":"2025-05-11T11:56:58.568872Z","iopub.status.idle":"2025-05-11T11:56:58.655926Z","shell.execute_reply.started":"2025-05-11T11:56:58.568850Z","shell.execute_reply":"2025-05-11T11:56:58.655307Z"},"papermill":{"duration":0.372005,"end_time":"2022-05-12T10:16:51.63863","exception":false,"start_time":"2022-05-12T10:16:51.266625","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"val_df[\"cell_type\"].value_counts()","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:56:58.656658Z","iopub.execute_input":"2025-05-11T11:56:58.656850Z","iopub.status.idle":"2025-05-11T11:56:58.666177Z","shell.execute_reply.started":"2025-05-11T11:56:58.656826Z","shell.execute_reply":"2025-05-11T11:56:58.665559Z"},"papermill":{"duration":0.27416,"end_time":"2022-05-12T10:16:52.16563","exception":false,"start_time":"2022-05-12T10:16:51.89147","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df_mark = train_df[train_df[\"cell_type\"] == \"markdown\"].reset_index(drop=True)\n\nval_df_mark = val_df[val_df[\"cell_type\"] == \"markdown\"].reset_index(drop=True)","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:56:58.667151Z","iopub.execute_input":"2025-05-11T11:56:58.667731Z","iopub.status.idle":"2025-05-11T11:56:58.762756Z","shell.execute_reply.started":"2025-05-11T11:56:58.667695Z","shell.execute_reply":"2025-05-11T11:56:58.762177Z"},"papermill":{"duration":0.371916,"end_time":"2022-05-12T10:16:52.797271","exception":false,"start_time":"2022-05-12T10:16:52.425355","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.metrics import mean_squared_error\n\nmean_squared_error(val_df_mark[\"pct_rank\"], np.ones(val_df_mark.shape[0])*train_df_mark[\"pct_rank\"].mean())","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:56:58.763648Z","iopub.execute_input":"2025-05-11T11:56:58.763885Z","iopub.status.idle":"2025-05-11T11:56:58.778827Z","shell.execute_reply.started":"2025-05-11T11:56:58.763855Z","shell.execute_reply":"2025-05-11T11:56:58.778177Z"},"papermill":{"duration":0.280778,"end_time":"2022-05-12T10:16:53.349877","exception":false,"start_time":"2022-05-12T10:16:53.069099","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from tqdm import tqdm\nimport sys, os\n\nfrom transformers import T5Tokenizer, T5ForConditionalGeneration\nimport torch.nn.functional as F\nimport torch.nn as nn\nimport torch\n\nMAX_LEN = 128\n\nclass MarkdownModel(nn.Module):\n    def __init__(self):\n        super(MarkdownModel, self).__init__()\n        \n        self.t5_model = T5ForConditionalGeneration.from_pretrained(T5_PATH)\n        self.tokenizer = T5Tokenizer.from_pretrained(T5_PATH)\n        self.top = nn.Linear(32128, 1)\n        \n    def forward(self, input_ids, attention_mask, target=None):\n        \n        outputs = self.t5_model(input_ids=input_ids, attention_mask=attention_mask, decoder_input_ids =target)[0]\n        \n       \n        return self.top(outputs[:, 0, :])","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:56:58.779625Z","iopub.execute_input":"2025-05-11T11:56:58.779829Z","iopub.status.idle":"2025-05-11T11:57:05.661095Z","shell.execute_reply.started":"2025-05-11T11:56:58.779808Z","shell.execute_reply":"2025-05-11T11:57:05.660383Z"},"papermill":{"duration":7.145711,"end_time":"2022-05-12T10:17:00.757077","exception":false,"start_time":"2022-05-12T10:16:53.611366","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch.utils.data import DataLoader, Dataset\n\n\n\nclass MarkdownDataset(Dataset):\n    \n    def __init__(self, df, max_len):\n        super().__init__()\n        self.df = df.reset_index(drop=True)\n        self.max_len = max_len\n        self.tokenizer = T5Tokenizer.from_pretrained(T5_PATH)\n\n    def __getitem__(self, index):\n        row = self.df.iloc[index]\n        \n        in_shot = \"Get this python notebook percentage rank: \\n \\n\" + row.source\n        #print(in_shot)\n        inputs = self.tokenizer.encode_plus(\n            in_shot,\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\n        target_shot = f\"This python notebook percentage rank is {row.pct_rank}\"\n        \n        decoder_inputs_aka_target = self.tokenizer.encode_plus(\n            target_shot,\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        \n        ids = torch.LongTensor(inputs['input_ids'])\n        mask = torch.LongTensor(inputs['attention_mask'])\n        d_ids = torch.LongTensor(decoder_inputs_aka_target['input_ids']) # maybe add attention mask for decoder\n\n        return ids, mask,d_ids, torch.FloatTensor([row.pct_rank])\n\n    def __len__(self):\n        return self.df.shape[0]\n    \ntrain_ds = MarkdownDataset(train_df_mark, max_len=MAX_LEN)\nval_ds = MarkdownDataset(val_df_mark, max_len=MAX_LEN)\n\nval_ds[0]","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:57:05.662294Z","iopub.execute_input":"2025-05-11T11:57:05.662858Z","iopub.status.idle":"2025-05-11T11:57:05.904840Z","shell.execute_reply.started":"2025-05-11T11:57:05.662821Z","shell.execute_reply":"2025-05-11T11:57:05.904184Z"},"papermill":{"duration":0.474499,"end_time":"2022-05-12T10:17:01.487031","exception":false,"start_time":"2022-05-12T10:17:01.012532","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"val_df_mark.iloc[0]","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:57:05.905772Z","iopub.execute_input":"2025-05-11T11:57:05.905965Z","iopub.status.idle":"2025-05-11T11:57:05.911684Z","shell.execute_reply.started":"2025-05-11T11:57:05.905942Z","shell.execute_reply":"2025-05-11T11:57:05.911048Z"},"papermill":{"duration":0.317331,"end_time":"2022-05-12T10:17:02.061139","exception":false,"start_time":"2022-05-12T10:17:01.743808","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"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),\n                                 eps=1e-08)\n    return optimizer","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:57:05.912548Z","iopub.execute_input":"2025-05-11T11:57:05.912765Z","iopub.status.idle":"2025-05-11T11:57:05.921447Z","shell.execute_reply.started":"2025-05-11T11:57:05.912740Z","shell.execute_reply":"2025-05-11T11:57:05.920897Z"},"papermill":{"duration":0.265988,"end_time":"2022-05-12T10:17:02.580374","exception":false,"start_time":"2022-05-12T10:17:02.314386","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BS = 32\nNW = 8\n\ntrain_loader = DataLoader(train_ds, batch_size=BS, shuffle=True, num_workers=NW,\n                          pin_memory=False, drop_last=True)\nval_loader = DataLoader(val_ds, batch_size=BS, shuffle=False, num_workers=NW,\n                          pin_memory=False, drop_last=False)","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:57:05.922231Z","iopub.execute_input":"2025-05-11T11:57:05.922402Z","iopub.status.idle":"2025-05-11T11:57:05.934607Z","shell.execute_reply.started":"2025-05-11T11:57:05.922381Z","shell.execute_reply":"2025-05-11T11:57:05.933946Z"},"papermill":{"duration":0.298424,"end_time":"2022-05-12T10:17:03.132","exception":false,"start_time":"2022-05-12T10:17:02.833576","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def read_data(data):\n    return tuple(d.cuda() for d in data[:-1]), data[-1].cuda()\n\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[0], inputs[1],inputs[2])\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 train(model, train_loader, val_loader, epochs):\n    np.random.seed(0)\n    \n    optimizer = get_optimizer(model)\n\n    criterion = torch.nn.MSELoss()\n    \n    for e in range(epochs):   \n        model.train()\n        tbar = tqdm(train_loader, file=sys.stdout)\n        \n        lr = adjust_lr(optimizer, e)\n        \n        loss_list = []\n        preds = []\n        labels = []\n\n        for idx, data in enumerate(tbar):\n            inputs, target = read_data(data)\n\n            optimizer.zero_grad()\n            pred = model(inputs[0], inputs[1],inputs[2])\n\n            loss = criterion(pred, target)\n            loss.backward()\n            optimizer.step()\n            \n            loss_list.append(loss.detach().cpu().item())\n            preds.append(pred.detach().cpu().numpy().ravel())\n            labels.append(target.detach().cpu().numpy().ravel())\n            \n            avg_loss = np.round(np.mean(loss_list), 4)\n\n            tbar.set_description(f\"Epoch {e+1} Loss: {avg_loss} lr: {lr}\")\n            \n        y_val, y_pred = validate(model, val_loader)\n            \n        print(\"Validation MSE:\", np.round(mean_squared_error(y_val, y_pred), 4))\n        print()\n    return model, y_pred\n\nmodel = MarkdownModel()\nmodel = model.cuda()\nmodel, y_pred = train(model, train_loader, val_loader, epochs=1)","metadata":{"execution":{"iopub.status.busy":"2025-05-11T11:57:05.935819Z","iopub.execute_input":"2025-05-11T11:57:05.936093Z","iopub.status.idle":"2025-05-11T14:13:45.818383Z","shell.execute_reply.started":"2025-05-11T11:57:05.936066Z","shell.execute_reply":"2025-05-11T14:13:45.817559Z"},"papermill":{"duration":987.160977,"end_time":"2022-05-12T10:33:30.548236","exception":false,"start_time":"2022-05-12T10:17:03.387259","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"val_df[\"pred\"] = val_df.groupby([\"id\", \"cell_type\"])[\"rank\"].rank(pct=True)\nval_df.loc[val_df[\"cell_type\"] == \"markdown\", \"pred\"] = y_pred","metadata":{"execution":{"iopub.status.busy":"2025-05-11T14:13:45.820167Z","iopub.execute_input":"2025-05-11T14:13:45.820408Z","iopub.status.idle":"2025-05-11T14:13:45.876889Z","shell.execute_reply.started":"2025-05-11T14:13:45.820376Z","shell.execute_reply":"2025-05-11T14:13:45.876158Z"},"papermill":{"duration":3.695855,"end_time":"2022-05-12T10:33:37.305989","exception":false,"start_time":"2022-05-12T10:33:33.610134","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"y_dummy = val_df.sort_values(\"pred\").groupby('id')['cell_id'].apply(list)\nkendall_tau(df_orders.loc[y_dummy.index], y_dummy)","metadata":{"execution":{"iopub.status.busy":"2025-05-11T14:13:45.877999Z","iopub.execute_input":"2025-05-11T14:13:45.878499Z","iopub.status.idle":"2025-05-11T14:13:46.015288Z","shell.execute_reply.started":"2025-05-11T14:13:45.878462Z","shell.execute_reply":"2025-05-11T14:13:46.014469Z"},"papermill":{"duration":3.230655,"end_time":"2022-05-12T10:33:43.545936","exception":false,"start_time":"2022-05-12T10:33:40.315281","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"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":"2025-05-11T14:13:46.016381Z","iopub.execute_input":"2025-05-11T14:13:46.016997Z","iopub.status.idle":"2025-05-11T14:13:46.072400Z","shell.execute_reply.started":"2025-05-11T14:13:46.016956Z","shell.execute_reply":"2025-05-11T14:13:46.071729Z"},"papermill":{"duration":3.156008,"end_time":"2022-05-12T10:33:49.98707","exception":false,"start_time":"2022-05-12T10:33:46.831062","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df[\"rank\"] = test_df.groupby([\"id\", \"cell_type\"]).cumcount()\ntest_df[\"pred\"] = test_df.groupby([\"id\", \"cell_type\"])[\"rank\"].rank(pct=True)","metadata":{"execution":{"iopub.status.busy":"2025-05-11T14:13:46.075381Z","iopub.execute_input":"2025-05-11T14:13:46.075621Z","iopub.status.idle":"2025-05-11T14:13:46.084818Z","shell.execute_reply.started":"2025-05-11T14:13:46.075599Z","shell.execute_reply":"2025-05-11T14:13:46.084145Z"},"papermill":{"duration":3.580422,"end_time":"2022-05-12T10:33:56.648552","exception":false,"start_time":"2022-05-12T10:33:53.06813","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df[\"pct_rank\"] = 0\ntest_ds = MarkdownDataset(test_df[test_df[\"cell_type\"] == \"markdown\"].reset_index(drop=True), max_len=MAX_LEN)\ntest_loader = DataLoader(test_ds, batch_size=BS, shuffle=False, num_workers=NW,\n                          pin_memory=False, drop_last=False)\n\nlen(test_ds), test_ds[0]","metadata":{"execution":{"iopub.status.busy":"2025-05-11T14:13:46.085645Z","iopub.execute_input":"2025-05-11T14:13:46.085882Z","iopub.status.idle":"2025-05-11T14:13:46.158282Z","shell.execute_reply.started":"2025-05-11T14:13:46.085851Z","shell.execute_reply":"2025-05-11T14:13:46.157585Z"},"papermill":{"duration":3.130783,"end_time":"2022-05-12T10:34:03.120038","exception":false,"start_time":"2022-05-12T10:33:59.989255","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"_, y_test = validate(model, test_loader)","metadata":{"execution":{"iopub.status.busy":"2025-05-11T14:13:46.159253Z","iopub.execute_input":"2025-05-11T14:13:46.159740Z","iopub.status.idle":"2025-05-11T14:13:46.821312Z","shell.execute_reply.started":"2025-05-11T14:13:46.159705Z","shell.execute_reply":"2025-05-11T14:13:46.820402Z"},"papermill":{"duration":4.00212,"end_time":"2022-05-12T10:34:10.223448","exception":false,"start_time":"2022-05-12T10:34:06.221328","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df.loc[test_df[\"cell_type\"] == \"markdown\", \"pred\"] = y_test","metadata":{"execution":{"iopub.status.busy":"2025-05-11T14:13:46.822994Z","iopub.execute_input":"2025-05-11T14:13:46.823635Z","iopub.status.idle":"2025-05-11T14:13:46.830232Z","shell.execute_reply.started":"2025-05-11T14:13:46.823587Z","shell.execute_reply":"2025-05-11T14:13:46.829580Z"},"papermill":{"duration":3.164567,"end_time":"2022-05-12T10:34:16.415308","exception":false,"start_time":"2022-05-12T10:34:13.250741","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"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.head()","metadata":{"execution":{"iopub.status.busy":"2025-05-11T14:13:46.831154Z","iopub.execute_input":"2025-05-11T14:13:46.831388Z","iopub.status.idle":"2025-05-11T14:13:46.849042Z","shell.execute_reply.started":"2025-05-11T14:13:46.831351Z","shell.execute_reply":"2025-05-11T14:13:46.848389Z"},"papermill":{"duration":3.093752,"end_time":"2022-05-12T10:34:22.827853","exception":false,"start_time":"2022-05-12T10:34:19.734101","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sub_df.to_csv(\"submission.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2025-05-11T14:13:46.849850Z","iopub.execute_input":"2025-05-11T14:13:46.850037Z","iopub.status.idle":"2025-05-11T14:13:46.863351Z","shell.execute_reply.started":"2025-05-11T14:13:46.850014Z","shell.execute_reply":"2025-05-11T14:13:46.862629Z"},"papermill":{"duration":3.551878,"end_time":"2022-05-12T10:34:29.528868","exception":false,"start_time":"2022-05-12T10:34:25.97699","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"papermill":{"duration":3.030555,"end_time":"2022-05-12T10:34:35.897014","exception":false,"start_time":"2022-05-12T10:34:32.866459","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null}]}