{"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":"markdown","source":"## 🧑‍💻 __AI4Code DeBERTa Train & Infer__\n\n---\n### <a href='#hyperparameters'> ⚙️ Hyperparameters </a> | <a href='#data-factory'> ⚒ Data Factory </a>  | <a href='#model'> 🧠 Model </a>  | <a href='#training'> ⚡ Training </a> \n","metadata":{"papermill":{"duration":0.010833,"end_time":"2022-07-02T01:42:25.719774","exception":false,"start_time":"2022-07-02T01:42:25.708941","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Sync Notebook with VS Code #\nimport sys\nsys.path.append('/kaggle/input/github-ai4code/fast-nlp')\nsys.path.append('/kaggle/input/github-ai4code/ai4code')\nsys.path.append('/kaggle/input/omegaconf')\nsys.path.extend(['fast-nlp', 'ai4code'])\nfrom src import *\n\nimport ai4c\nimport ai4c.process_df\nfrom ai4c.submission import CELL_SEP, IS_INTERACTIVE\n\n!touch __init__.py","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:42:25.744751Z","iopub.status.busy":"2022-07-02T01:42:25.743663Z","iopub.status.idle":"2022-07-02T01:43:07.183769Z","shell.execute_reply":"2022-07-02T01:43:07.182732Z"},"papermill":{"duration":41.454329,"end_time":"2022-07-02T01:43:07.18597","exception":false,"start_time":"2022-07-02T01:42:25.731641","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# ⚙️ Hyperparameters ⚙️\n---\n### <a href='#data-factory'> ⚒ Data Factory </a>  | <a href='#model'> 🧠 Model </a>|  <a href='#training'> ⚡ Training </a> \n\n<a name='hyperparameters'>","metadata":{"papermill":{"duration":0.005365,"end_time":"2022-07-02T01:43:07.197104","exception":false,"start_time":"2022-07-02T01:43:07.191739","status":"completed"},"tags":[]}},{"cell_type":"code","source":"%%hyperparameters\n\n## Huggingface Backbone ##\nbackbone_name: 'microsoft/deberta-v3-large'\nbackbone_dir: 'deberta-backbones'\n\nattention_probs_dropout_prob: 0.05\nhidden_dropout_prob: 0.10\n\nmax_relative_positions: 512\nmax_seq_len: 1280\n\n## Data Factory ##\nmax_markdown_seq_len: 1024\ntrain_folds: [1, 2]\nvalid_fold: 0\nmax_tokens_per_cell: 256\n\n## Model Training ##\nnum_train_epochs: 1\ntrain_batch_size: 1\neval_batch_size: 4\n\ngradient_accumulation_steps: 16\nmixed_precision: True\n\n## Loss Function ##\nmarkdown_cell_loss_weight: 0.50\n\n## Cosine Decay LR Scheduler ##\nwarmup_ratio: 0.125\nlearning_rate: 1e-4\n\n## AdamW Optimizer ##\nweight_decay: 1e-3\nmax_grad_norm: 1000000.0\nadam_epsilon: 1e-6\n\n## Logging ##\nlogging_freq: 100\ncheckpoint_freq: 1000\n\n## Inference ##\nhide_lb_score: False","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:43:07.209885Z","iopub.status.busy":"2022-07-02T01:43:07.209019Z","iopub.status.idle":"2022-07-02T01:43:07.300859Z","shell.execute_reply":"2022-07-02T01:43:07.30018Z"},"papermill":{"duration":0.099943,"end_time":"2022-07-02T01:43:07.302429","exception":false,"start_time":"2022-07-02T01:43:07.202486","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"backbone_code = 'deberta_v3_large'\nbackbone_dir = f'/kaggle/input/{HP.backbone_dir}/{backbone_code}'\n\nbackbone = AutoModel.from_pretrained(\n    backbone_dir,\n    attention_probs_dropout_prob=HP.attention_probs_dropout_prob,\n    hidden_dropout_prob=HP.hidden_dropout_prob,\n    # max_relative_positions=HP.max_relative_positions,\n    type_vocab_size=2,\n)\ntokenizer = AutoTokenizer.from_pretrained(backbone_dir, use_fast=True)","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:43:07.314201Z","iopub.status.busy":"2022-07-02T01:43:07.313949Z","iopub.status.idle":"2022-07-02T01:43:27.33969Z","shell.execute_reply":"2022-07-02T01:43:27.338887Z"},"papermill":{"duration":20.034363,"end_time":"2022-07-02T01:43:27.342145","exception":false,"start_time":"2022-07-02T01:43:07.307782","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"notebooks_df = pd.read_csv('/kaggle/input/ai4code-dataframes/notebooks_df.csv')\ntrain_df = notebooks_df[notebooks_df.notebook_fold.isin(HP.train_folds)]\nvalid_df = notebooks_df[notebooks_df.notebook_fold == HP.valid_fold]\nif Path('/kaggle/input/AI4Code').exists():\n    cell_df_test = ai4c.process_df.build_cell_df('/kaggle/input/AI4Code/test')\n    test_df = ai4c.process_df.build_notebooks_df(cell_df_test)\nelse: \n    test_df = valid_df.drop(columns=['merged_cell_pct_ranks'])\n\nvalid_df = valid_df.sample(100)\nif len(test_df) < 100:\n    train_df = train_df.sample(300)","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:43:27.355367Z","iopub.status.busy":"2022-07-02T01:43:27.3549Z","iopub.status.idle":"2022-07-02T01:44:46.424714Z","shell.execute_reply":"2022-07-02T01:44:46.423851Z"},"papermill":{"duration":79.078194,"end_time":"2022-07-02T01:44:46.426539","exception":false,"start_time":"2022-07-02T01:43:27.348345","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%writefile data_module.py\n\nfrom collections import defaultdict\nfrom tqdm.auto import tqdm \nfrom pathlib import Path\nfrom time import time\nimport pandas as pd\n\nimport transformers\nimport datasets \n\nimport torch\nimport ai4c\n\nCELL_SEP = '[CELL_SEP]'\ntqdm.pandas()\n\ndef prune_cell_tokens(cell_token_ids, max_seq_len):\n    cell_token_counts = [len(token_ids) for token_ids in cell_token_ids]\n    total_number_of_cells = len(cell_token_counts)\n    total_tokens_to_prune = max(sum(cell_token_counts)-max_seq_len, 0)\n\n    tokens_to_prune_per_cell = [0]*total_number_of_cells\n    total_pruned_tokens = 0\n    while total_tokens_to_prune > 0:\n        cur_max_cell_token_count = max(cell_token_counts)\n        second_max_cell_token_count = sorted(cell_token_counts)[-2]\n        for cell_idx, cell_token_count in enumerate(cell_token_counts):\n            if not cell_token_count == cur_max_cell_token_count: \n                continue\n            \n            num_tokens_to_pop = min(cell_token_count-second_max_cell_token_count+1, total_tokens_to_prune)\n            tokens_to_prune_per_cell[cell_idx] += num_tokens_to_pop\n            total_pruned_tokens += num_tokens_to_pop\n            total_tokens_to_prune -= num_tokens_to_pop\n            cell_token_counts[cell_idx] -= num_tokens_to_pop\n            break\n    \n    # Prune the cell tokens\n    pruned_cell_token_ids = []\n    for cell_token_ids, num_tokens_to_pop in zip(cell_token_ids, tokens_to_prune_per_cell):\n        if num_tokens_to_pop == 0:\n            pruned_cell_token_ids.append(cell_token_ids)\n            continue\n        pruned_cell_token_ids.append(cell_token_ids[:-num_tokens_to_pop])\n    return pruned_cell_token_ids\n\n\ndef convert_to_features_torch_rankencoder(notebook, tokenizer, max_seq_len, max_markdown_seq_len, max_tokens_per_cell): \n    '''Tokenizer the notebook and convert to features for the model.'''\n\n    markdown_cell_sources = notebook['merged_markdown_cell_sources'].split(CELL_SEP)\n    code_cell_sources = notebook['merged_code_cell_sources'].split(CELL_SEP)\n    \n    if 'merged_markdown_cell_pct_ranks' in notebook:\n        markdown_cell_pct_ranks = [float(rank) for rank in notebook['merged_markdown_cell_pct_ranks'].split(CELL_SEP)]\n        code_cell_pct_ranks = [float(rank) for rank in notebook['merged_code_cell_pct_ranks'].split(CELL_SEP)]\n    else:\n        markdown_cell_pct_ranks = [0]*len(markdown_cell_sources)\n        code_cell_pct_ranks = [0]*len(code_cell_sources)\n\n    # Remove cells from the end of the notebook so that all cells have at least one representative token\n    max_markdown_cells = max_markdown_seq_len//2\n    max_code_cells = (max_seq_len-max_markdown_seq_len)//2\n    if len(markdown_cell_sources) > max_markdown_cells:\n        markdown_cell_sources = markdown_cell_sources[:max_markdown_cells]\n        markdown_cell_pct_ranks = markdown_cell_pct_ranks[:max_markdown_cells]\n    if len(code_cell_sources) > max_code_cells:\n        code_cell_sources = code_cell_sources[:max_code_cells]\n        code_cell_pct_ranks = code_cell_pct_ranks[:max_code_cells]\n    markdown_cell_count = len(markdown_cell_sources)\n    code_cell_count = len(code_cell_sources)\n\n    max_tokens_per_markdown_cell = max(max_tokens_per_cell, max_markdown_seq_len//markdown_cell_count)\n    markdown_cell_token_ids = tokenizer(\n        markdown_cell_sources,\n        max_length=max_tokens_per_markdown_cell,\n        truncation=True,\n    )['input_ids']\n    markdown_cell_token_ids = prune_cell_tokens(markdown_cell_token_ids, max_markdown_seq_len)\n    total_markdown_cell_tokens = sum([len(token_ids) for token_ids in markdown_cell_token_ids])\n\n    max_code_seq_len = max_seq_len - total_markdown_cell_tokens\n    max_tokens_per_code_cell = max(max_tokens_per_cell, max_code_seq_len//code_cell_count)\n    code_cell_token_ids = tokenizer(\n        code_cell_sources, \n        max_length=max_tokens_per_code_cell, \n        truncation=True, \n    )['input_ids']\n    code_cell_token_ids = prune_cell_tokens(code_cell_token_ids, max_seq_len-total_markdown_cell_tokens)\n\n    # Merge the tokenized cells and create the model features\n    cell_token_ids = markdown_cell_token_ids + code_cell_token_ids\n    notebook_cell_count = len(cell_token_ids)\n\n    # Create the model features\n    if 'merged_cell_pct_ranks' in notebook:\n        cell_pct_ranks = markdown_cell_pct_ranks + code_cell_pct_ranks\n    else: \n        cell_pct_ranks = [None]*notebook_cell_count\n    \n    input_ids = []\n    token_cell_ids, token_type_ids = [], []\n    token_weights, token_labels = [], []\n    for cur_cell_idx, cell_token_ids in enumerate(cell_token_ids):\n        token_count_for_cell = len(cell_token_ids)\n        if cur_cell_idx < markdown_cell_count:\n            token_cell_id = 1.0\n        else: \n            token_cell_id = 2.0\n        \n        input_ids += cell_token_ids\n        token_cell_ids += [token_cell_id] * token_count_for_cell\n        token_type_id = 0 if cur_cell_idx % 2 == 0 else 1\n        token_type_ids += [token_type_id] * token_count_for_cell\n\n        token_labels += [cell_pct_ranks[cur_cell_idx]] * token_count_for_cell\n        token_weights += [1/token_count_for_cell] * token_count_for_cell\n    \n    # Build the feature dict for the input \n    notebook_features = {\n        'input_ids': input_ids, \n        'token_type_ids': token_type_ids,\n        'token_cell_ids': token_cell_ids,\n    }\n    if 'merged_cell_pct_ranks' in notebook:\n        notebook_features['token_labels'] = token_labels\n        notebook_features['token_weights'] = token_weights\n    return notebook_features\n\nclass AI4CodeDataset(torch.utils.data.Dataset):\n    def __init__(\n        self,\n        df,\n        tokenizer,\n        max_seq_len,\n        max_markdown_seq_len,\n        max_tokens_per_cell,\n    ):\n        self.df = df\n        self.tokenizer = tokenizer\n        self.max_seq_len = max_seq_len\n        self.max_markdown_seq_len = max_markdown_seq_len\n        self.max_tokens_per_cell = max_tokens_per_cell\n        self.is_test = 'merged_cell_pct_ranks' not in df.columns\n    \n    def __getitem__(self, idx):\n        notebook_row = self.df.iloc[idx]\n        notebook_features = convert_to_features_torch_rankencoder(\n            notebook=notebook_row,\n            tokenizer=self.tokenizer,\n            max_seq_len=self.max_seq_len,\n            max_markdown_seq_len=self.max_markdown_seq_len,\n            max_tokens_per_cell=self.max_tokens_per_cell,\n        )\n        inputs = {\n            'input_ids': notebook_features['input_ids'],\n            'token_type_ids': notebook_features['token_type_ids'],\n            'token_cell_ids': notebook_features['token_cell_ids'],\n        }\n        if self.is_test:\n            return inputs\n        \n        labels = {\n            'token_labels': notebook_features['token_labels'],\n            'token_weights': notebook_features['token_weights'],\n        }\n        return inputs, labels\n\n    def __len__(self):\n        return len(self.df)\n    \n    \ndef train_collate_fn(batch, tokenizer):\n    batch_input_ids = [{'input_ids': inputs['input_ids']} for inputs, labels in batch]\n    padded_batch = tokenizer.pad(\n        batch_input_ids,\n        pad_to_multiple_of=8,\n        padding=True, \n    )\n    feature_to_pad_token_id = {\n        'token_type_ids': -100.0,\n        'token_cell_ids': -100.0,\n        'token_labels': -100.0,\n        'token_weights': 0.0,\n    }\n    \n    batch_inputs = {\n        'input_ids': padded_batch['input_ids'],\n        'attention_mask': padded_batch['attention_mask'],\n        'token_type_ids': [],\n        'token_cell_ids': [],\n    }\n    batch_labels = {\n        'token_weights': [], \n        'token_labels': [],\n    }\n    \n    batch_sequence_length = torch.tensor(padded_batch['input_ids']).shape[1]\n    for example_inputs, example_labels in batch:\n        for feature, value in example_inputs.items():\n            if feature not in feature_to_pad_token_id:\n                continue\n            num_tokens_to_pad = batch_sequence_length-len(value)\n            pad_token_id = feature_to_pad_token_id[feature]\n            batch_inputs[feature].append(list(value) + [pad_token_id]*num_tokens_to_pad)\n        \n        for feature, value in example_labels.items():\n            if feature not in feature_to_pad_token_id:\n                continue\n            num_tokens_to_pad = batch_sequence_length-len(value)\n            pad_token_id = feature_to_pad_token_id[feature]\n            batch_labels[feature].append(list(value) + [pad_token_id]*num_tokens_to_pad)\n    \n    batch_inputs = {\n        'input_ids': torch.tensor(batch_inputs['input_ids'], dtype=torch.long),\n        'attention_mask': torch.tensor(batch_inputs['attention_mask'], dtype=torch.long),\n        'token_type_ids': torch.tensor(batch_inputs['token_type_ids'], dtype=torch.long),\n        'token_cell_ids': torch.tensor(batch_inputs['token_cell_ids'], dtype=torch.long),\n    }\n    batch_labels = {\n        'token_weights': torch.tensor(batch_labels['token_weights'], dtype=torch.float),\n        'token_labels': torch.tensor(batch_labels['token_labels'], dtype=torch.float),\n    }\n    return batch_inputs, batch_labels\n\n\ndef test_collate_fn(batch, tokenizer):\n    batch_input_ids = [{'input_ids': inputs['input_ids']} for inputs in batch]\n    padded_batch = tokenizer.pad(\n        batch_input_ids,\n        pad_to_multiple_of=8,\n        padding=True, \n    )\n    feature_to_pad_token_id = {\n        'token_type_ids': -100.0,\n        'token_cell_ids': -100.0,\n        'token_weights': 0.0,\n    }\n    batch_inputs = {\n        'input_ids': padded_batch['input_ids'],\n        'attention_mask': padded_batch['attention_mask'],\n        'token_type_ids': [],\n        'token_cell_ids': [],\n    }\n    \n    batch_sequence_length = torch.tensor(padded_batch['input_ids']).shape[1]\n    for example_inputs in batch:\n        for feature, value in example_inputs.items():\n            if feature not in feature_to_pad_token_id:\n                continue\n            num_tokens_to_pad = batch_sequence_length-len(value)\n            pad_token_id = feature_to_pad_token_id[feature]\n            batch_inputs[feature].append(list(value) + [pad_token_id]*num_tokens_to_pad)\n    \n    batch_inputs = {\n        'input_ids': torch.tensor(batch_inputs['input_ids'], dtype=torch.long),\n        'attention_mask': torch.tensor(batch_inputs['attention_mask'], dtype=torch.long),\n        'token_type_ids': torch.tensor(batch_inputs['token_type_ids'], dtype=torch.long),\n        'token_cell_ids': torch.tensor(batch_inputs['token_cell_ids'], dtype=torch.long),\n    }\n    return batch_inputs","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:44:46.440773Z","iopub.status.busy":"2022-07-02T01:44:46.440252Z","iopub.status.idle":"2022-07-02T01:44:46.552628Z","shell.execute_reply":"2022-07-02T01:44:46.550915Z"},"papermill":{"duration":0.12194,"end_time":"2022-07-02T01:44:46.554652","exception":false,"start_time":"2022-07-02T01:44:46.432712","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import data_module\n\ntrain_dataset = data_module.AI4CodeDataset(\n    df=train_df,\n    tokenizer=tokenizer,\n    max_seq_len=HP.max_seq_len,\n    max_markdown_seq_len=HP.max_markdown_seq_len,\n    max_tokens_per_cell=HP.max_tokens_per_cell,\n)\n\neval_dataset = data_module.AI4CodeDataset(\n    df=valid_df,\n    tokenizer=tokenizer,\n    max_seq_len=HP.max_seq_len,\n    max_markdown_seq_len=HP.max_markdown_seq_len,\n    max_tokens_per_cell=HP.max_tokens_per_cell,\n)\n\ntrain_collate_fn = partial(data_module.train_collate_fn, tokenizer=tokenizer)\ntrain_dataloader = torch.utils.data.DataLoader(\n    dataset=train_dataset,\n    batch_size=HP.train_batch_size,\n    shuffle=True,\n    num_workers=2,\n    collate_fn=train_collate_fn,\n    pin_memory=True,\n    drop_last=True,\n    prefetch_factor=4,\n)\n\neval_dataloader = torch.utils.data.DataLoader(\n    dataset=eval_dataset,\n    batch_size=HP.eval_batch_size,\n    shuffle=False,\n    num_workers=2,\n    collate_fn=train_collate_fn,\n    pin_memory=True,\n    drop_last=False,\n    prefetch_factor=4,\n)\n\nbatch = next(iter(train_dataloader))","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:44:46.569123Z","iopub.status.busy":"2022-07-02T01:44:46.568466Z","iopub.status.idle":"2022-07-02T01:44:54.182225Z","shell.execute_reply":"2022-07-02T01:44:54.181238Z"},"papermill":{"duration":7.623763,"end_time":"2022-07-02T01:44:54.184682","exception":false,"start_time":"2022-07-02T01:44:46.560919","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%writefile torch_model.py\n\nimport transformers\nimport torch.nn as nn\nimport torch\n\nclass AI4CodeModel(nn.Module):\n    def __init__(self, backbone):\n        super().__init__()\n        self.backbone = backbone\n        self.ranker = nn.Linear(backbone.config.hidden_size, 1)\n    \n    def forward(self, input_ids, attention_mask, token_type_ids):\n        backbone_outputs = self.backbone(\n            input_ids=input_ids,\n            # attention_mask=attention_mask,\n            # token_type_ids=token_type_ids,\n        )\n        token_preds = self.ranker(backbone_outputs.last_hidden_state)\n        seq_len = token_preds.size(1)\n        return token_preds.view((-1, seq_len))\n\ndef get_optimizer_grouped_parameters(model, weight_decay):\n    no_decay = [\"bias\", \"LayerNorm.weight\"]\n    optimizer_grouped_parameters = [\n        {\n            \"params\": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],\n            \"weight_decay\": weight_decay,\n        },\n        {\n            \"params\": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],\n            \"weight_decay\": 0.0,\n        },\n    ]\n    return optimizer_grouped_parameters","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:44:54.198545Z","iopub.status.busy":"2022-07-02T01:44:54.19822Z","iopub.status.idle":"2022-07-02T01:44:54.283974Z","shell.execute_reply":"2022-07-02T01:44:54.28326Z"},"papermill":{"duration":0.09485,"end_time":"2022-07-02T01:44:54.285795","exception":false,"start_time":"2022-07-02T01:44:54.190945","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch_model\n\nmodel = torch_model.AI4CodeModel(backbone)\noptimizer = torch.optim.AdamW(\n    params=torch_model.get_optimizer_grouped_parameters(model.backbone, HP.weight_decay), \n    eps=HP.adam_epsilon,\n    lr=HP.learning_rate,\n)\nnum_update_steps_per_epoch = math.ceil(len(train_dataloader) / HP.gradient_accumulation_steps)\ntotal_train_steps = HP.num_train_epochs * num_update_steps_per_epoch\n\nlr_scheduler = transformers.get_cosine_schedule_with_warmup(\n    optimizer=optimizer,\n    num_warmup_steps=int(HP.warmup_ratio*total_train_steps),\n    num_training_steps=total_train_steps,\n)\n\n_ = model.cuda()\ntorch.cuda.empty_cache()\ngc.collect()","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:44:54.300303Z","iopub.status.busy":"2022-07-02T01:44:54.299633Z","iopub.status.idle":"2022-07-02T01:44:55.121445Z","shell.execute_reply":"2022-07-02T01:44:55.120633Z"},"papermill":{"duration":0.831065,"end_time":"2022-07-02T01:44:55.123266","exception":false,"start_time":"2022-07-02T01:44:54.292201","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# ⚡ Training ⚡\n---\n### <a href='#hyperparameters'> ⚙️ Hyperparameters </a>  | <a href='#model'> 🧠 Model </a>\n\n<a name='training'>","metadata":{"papermill":{"duration":0.006112,"end_time":"2022-07-02T01:44:55.135616","exception":false,"start_time":"2022-07-02T01:44:55.129504","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def mse_loss(token_labels, token_preds, token_weights, mask):\n    \"\"\"\n    Compute the mean squared error loss.\n    Args:\n        token_labels: tensor of shape (batch_size, seq_len)\n        token_preds: tensor of shape (batch_size, seq_len)\n        token_weights: tensor of shape (batch_size, seq_len)\n        mask: tensor of shape (batch_size, seq_len)\n    Returns:\n        mse_loss: scalar\n    \"\"\"\n    token_labels, token_preds, token_weights = token_labels*mask, token_preds*mask, token_weights*mask\n    sum_weights = torch.sum(token_weights, dim=-1)\n    notebook_losses = torch.sum(((token_labels - token_preds)**2 * token_weights), dim=-1) / (sum_weights + 1e-8)\n    return torch.mean(notebook_losses)\n\nprogress_bar = tqdm(range(total_train_steps), desc=\"Training Progress\") \ncompleted_batches = 0\n_ = model.train()\nscaler = torch.cuda.amp.GradScaler(enabled=HP.mixed_precision)\n\nepoch_metrics = {\n    'loss': 0.0, \n    'markdown_cell_loss': 0.0, \n    'code_cell_loss': 0.0\n}\nbatch_metrics = {\n    'loss': 0.0,\n    'markdown_cell_loss': 0.0,\n    'code_cell_loss': 0.0,\n    'num_markdown_cell_tokens': 0.0, \n    'num_code_cell_tokens': 0.0,\n    'num_pad_tokens': 0.0,\n}\n\nfor step, batch in enumerate(train_dataloader):\n    inputs, labels = batch\n    for k, v in inputs.items(): \n        inputs[k] = v.cuda()\n    for k, v in labels.items():\n        labels[k] = v.cuda()\n    \n    with torch.cuda.amp.autocast(enabled=HP.mixed_precision):\n        token_preds = model(\n            input_ids=inputs['input_ids'], \n            attention_mask=inputs['attention_mask'],\n            token_type_ids=inputs['token_type_ids'],\n        )\n        \n        markdown_cell_mask = torch.where(inputs['token_cell_ids']==1.0, 1.0, 0.0)\n        markdown_cell_loss = mse_loss(\n            token_labels=labels['token_labels'],\n            token_preds=token_preds,\n            token_weights=labels['token_weights'],\n            mask=markdown_cell_mask,\n        )\n\n        code_cell_mask = torch.where(inputs['token_cell_ids']==2.0, 1.0, 0.0)\n        code_cell_loss = mse_loss(\n            token_labels=labels['token_labels'],\n            token_preds=token_preds,\n            token_weights=labels['token_weights'],\n            mask=code_cell_mask,\n        )\n\n        loss = HP.markdown_cell_loss_weight*markdown_cell_loss + (1-HP.markdown_cell_loss_weight) * code_cell_loss\n    \n    epoch_metrics['loss'] += loss.item()\n    epoch_metrics['markdown_cell_loss'] += markdown_cell_loss.item()\n    epoch_metrics['code_cell_loss'] += code_cell_loss.item()\n\n    batch_metrics['loss'] += loss.item()\n    batch_metrics['markdown_cell_loss'] += markdown_cell_loss.item()\n    batch_metrics['code_cell_loss'] += code_cell_loss.item()\n    batch_metrics['num_markdown_cell_tokens'] += torch.mean(torch.sum(markdown_cell_mask, dim=-1)).item()\n    batch_metrics['num_code_cell_tokens'] += torch.mean(torch.sum(code_cell_mask, dim=-1)).item()\n    batch_metrics['num_pad_tokens'] += torch.mean(torch.sum(1.0-inputs['attention_mask'], dim=-1)).item()\n\n    scaler.scale(loss).backward()\n    if (step+1) % HP.gradient_accumulation_steps == 0:\n        grad_norm = torch.nn.utils.clip_grad_norm_(model.backbone.parameters(), HP.max_grad_norm)\n        scaler.step(optimizer)\n        scaler.update()\n        lr_scheduler.step()\n        optimizer.zero_grad()\n        completed_batches += 1\n        \n        progress_bar.set_postfix(\n            loss=f\"{batch_metrics['loss']/HP.gradient_accumulation_steps:.4f}\", \n            markdown_cell_loss=f\"{batch_metrics['markdown_cell_loss']/HP.gradient_accumulation_steps:.4f}\", \n            code_cell_loss=f\"{batch_metrics['code_cell_loss']/HP.gradient_accumulation_steps:.4f}\",\n            grad_norm=f\"{grad_norm:.4f}\",\n        )\n        _ = progress_bar.update(1)\n\n        if completed_batches % HP.logging_freq == 0:\n            print(f\"Batch #{completed_batches} | Pred shape: {token_preds.shape}\")\n            print(f\"Loss: {loss.item()} | Markdown Cell Loss: {markdown_cell_loss.item()} | Code Cell Loss: {code_cell_loss.item()}\")\n            print(f\"Gradient Norm: {grad_norm.item()}\")\n            print(f\"Learning Rate: {lr_scheduler.get_lr()[0]}\")\n            print(\"-\"*50)\n            batch_metrics = {key: 0 for key in batch_metrics}\n\nfor k, v in epoch_metrics.items():\n    print(f\"{colored(k, 'blue')}: {colored(v/(step+1), 'red')}\")","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:44:55.149158Z","iopub.status.busy":"2022-07-02T01:44:55.148593Z","iopub.status.idle":"2022-07-02T01:49:30.030082Z","shell.execute_reply":"2022-07-02T01:49:30.029082Z"},"papermill":{"duration":274.897068,"end_time":"2022-07-02T01:49:30.038567","exception":false,"start_time":"2022-07-02T01:44:55.141499","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 🛒 Validation\n---\n### <a href='#hyperparameters'> ⚙️ Hyperparameters </a>  | <a href='#model'> 🧠 Model </a>\n\n<a name='validation'>","metadata":{"papermill":{"duration":0.006402,"end_time":"2022-07-02T01:49:30.051854","exception":false,"start_time":"2022-07-02T01:49:30.045452","status":"completed"},"tags":[]}},{"cell_type":"code","source":"torch.cuda.empty_cache()\ngc.collect()\n_ = model.eval()\n\nepoch_metrics = {\n    'loss': 0.0, \n    'markdown_cell_loss': 0.0, \n    'code_cell_loss': 0.0\n}\ncell_id_to_rank_preds = defaultdict(list)\nnotebook_idx = 0\nfor step, (inputs, labels) in tqdm(enumerate(eval_dataloader), total=len(eval_dataloader)):\n    for k, v in inputs.items():\n        inputs[k] = v.cuda()\n    for k, v in labels.items():\n        labels[k] = v.cuda()\n\n    with torch.cuda.amp.autocast(enabled=HP.mixed_precision):\n        with torch.no_grad():\n            token_preds = model(\n                input_ids=inputs['input_ids'],\n                attention_mask=inputs['attention_mask'],\n                token_type_ids=inputs['token_type_ids'],\n            )\n            markdown_cell_mask = torch.where(inputs['token_cell_ids']==1.0, 1.0, 0.0)\n            markdown_cell_loss = mse_loss(\n                token_labels=labels['token_labels'],\n                token_preds=token_preds,\n                token_weights=labels['token_weights'],\n                mask=markdown_cell_mask,\n            )\n\n            code_cell_mask = torch.where(inputs['token_cell_ids']==2.0, 1.0, 0.0)\n            code_cell_loss = mse_loss(\n                token_labels=labels['token_labels'],\n                token_preds=token_preds,\n                token_weights=labels['token_weights'],\n                mask=code_cell_mask,\n            )\n            loss = HP.markdown_cell_loss_weight*markdown_cell_loss + (1-HP.markdown_cell_loss_weight) * code_cell_loss\n            \n            epoch_metrics['loss'] += loss.item()\n            epoch_metrics['markdown_cell_loss'] += markdown_cell_loss.item()\n            epoch_metrics['code_cell_loss'] += code_cell_loss.item()\n\n            token_preds = token_preds.detach().cpu().numpy()\n        \n        batch_size = inputs['input_ids'].shape[0]\n        batch_notebook_cell_ids = valid_df.iloc[notebook_idx:notebook_idx+batch_size].merged_cell_ids.values\n        for i in range(batch_size): \n            notebook_token_type_ids = inputs['token_type_ids'][i]\n            notebook_cell_ids = valid_df.iloc[notebook_idx+i].merged_cell_ids.split(CELL_SEP)\n            notebook_token_preds = token_preds[i]\n\n            cur_cell_idx = 0\n            prev_token_type_id = 0\n            for token_pred, token_type_id in zip(notebook_token_preds, notebook_token_type_ids):\n                if token_type_id == -100:\n                    continue\n\n                if token_type_id == prev_token_type_id:\n                    cur_cell_id = notebook_cell_ids[cur_cell_idx]\n                    cell_id_to_rank_preds[cur_cell_id].append(token_pred)\n                else: \n                    cur_cell_idx += 1\n                prev_token_type_id = token_type_id\n        notebook_idx += batch_size\n\nfor k, v in epoch_metrics.items():\n    print(f\"{colored(k, 'blue')}: {colored(v/(step+1), 'red')}\")\n\n# Compute Kendall Tau for the predictions with the ground truth\ncell_id_to_pred_rank = {cell_id: float(sum(preds)/len(preds)) for cell_id, preds in cell_id_to_rank_preds.items()}\nall_notebook_cell_pct_ranks = valid_df.merged_cell_pct_ranks.values\nall_notebook_cell_ids = valid_df.merged_cell_ids.values\nall_notebook_kendall_taus, all_notebook_cell_order_preds = [], []\nfor notebook_idx in range(len(valid_df)):\n    true_cell_ranks = [float(rank) for rank in all_notebook_cell_pct_ranks[notebook_idx].split(CELL_SEP)]\n    cell_ids = all_notebook_cell_ids[notebook_idx].split(CELL_SEP)\n    pred_cell_ranks = [cell_id_to_pred_rank.get(cell_id, cell_idx/len(cell_ids)) for cell_idx, cell_id in enumerate(cell_ids)]\n\n    notebook_tau = scipy.stats.kendalltau(true_cell_ranks, pred_cell_ranks, method='asymptotic')[0]\n    notebook_preds = CELL_SEP.join([str(round(rank, 4)) for rank in pred_cell_ranks])\n    all_notebook_kendall_taus.append(notebook_tau)\n    all_notebook_cell_order_preds.append(notebook_preds)\nall_notebook_kendall_taus = np.array(all_notebook_kendall_taus)\n\nvalid_df['kendall_tau'] = all_notebook_kendall_taus \navg_tau = all_notebook_kendall_taus.mean()\nprint(f\"{colored('Kendall Tau', 'blue')}: {colored(avg_tau, 'red')}\")\n\nfor cutoff in [4, 16, 64]:\n    tau = valid_df[valid_df.markdown_cell_count>cutoff].kendall_tau.mean()\n    print(f\"Kendall Tau for notebooks with {colored(cutoff, 'yellow')}+ markdown cells:\", colored(tau, 'red'))\n\nvalid_df.to_csv(f'valid_df.csv', index=False)","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:49:30.066496Z","iopub.status.busy":"2022-07-02T01:49:30.066168Z","iopub.status.idle":"2022-07-02T01:50:00.513109Z","shell.execute_reply":"2022-07-02T01:50:00.511722Z"},"papermill":{"duration":30.456811,"end_time":"2022-07-02T01:50:00.51501","exception":false,"start_time":"2022-07-02T01:49:30.058199","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tau_32 = valid_df[valid_df.markdown_cell_count>32].kendall_tau.mean()\nmodel_save_path = f\"{backbone_code}_tau{int(avg_tau*10000)}_tau32_{int(tau_32*1000)}.pt\"\ntorch.save(model.state_dict(), model_save_path)\nprint(f'Model saved at {model_save_path}')","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:50:00.530931Z","iopub.status.busy":"2022-07-02T01:50:00.530254Z","iopub.status.idle":"2022-07-02T01:50:04.765326Z","shell.execute_reply":"2022-07-02T01:50:04.764467Z"},"papermill":{"duration":4.245242,"end_time":"2022-07-02T01:50:04.76722","exception":false,"start_time":"2022-07-02T01:50:00.521978","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 🎯 Inference\n---\n### <a href='#hyperparameters'> ⚙️ Hyperparameters </a>  | <a href='#training'> ⚡ Training </a>\n\n<a name='inference'>","metadata":{"papermill":{"duration":0.009133,"end_time":"2022-07-02T01:50:04.78668","exception":false,"start_time":"2022-07-02T01:50:04.777547","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import data_module\ntest_dataset = data_module.AI4CodeDataset(\n    df=test_df,\n    tokenizer=tokenizer,\n    max_seq_len=HP.max_seq_len,\n    max_markdown_seq_len=HP.max_markdown_seq_len,\n    max_tokens_per_cell=HP.max_tokens_per_cell,\n)\n\ntest_collate_fn = partial(data_module.test_collate_fn, tokenizer=tokenizer)\ntest_dataloader = torch.utils.data.DataLoader(\n    dataset=test_dataset,\n    batch_size=HP.eval_batch_size,\n    shuffle=False,\n    num_workers=2,\n    collate_fn=test_collate_fn,\n    pin_memory=True,\n    drop_last=False,\n    prefetch_factor=4,\n)","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:50:04.9608Z","iopub.status.busy":"2022-07-02T01:50:04.960095Z","iopub.status.idle":"2022-07-02T01:50:05.037965Z","shell.execute_reply":"2022-07-02T01:50:05.037004Z"},"papermill":{"duration":0.089612,"end_time":"2022-07-02T01:50:05.040421","exception":false,"start_time":"2022-07-02T01:50:04.950809","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.cuda.empty_cache()\ngc.collect()\n\n_ = model.eval()\ncell_id_to_rank_preds = defaultdict(list)\nnotebook_idx = 0\nfor step, inputs in tqdm(enumerate(test_dataloader), total=len(test_dataloader)):\n    for k, v in inputs.items():\n        inputs[k] = v.cuda()\n    \n    with torch.cuda.amp.autocast(enabled=HP.mixed_precision):\n        with torch.no_grad():\n            token_preds = model(\n                input_ids=inputs['input_ids'],\n                attention_mask=inputs['attention_mask'],\n                token_type_ids=inputs['token_type_ids'],\n            )\n            token_preds = token_preds.detach().cpu().numpy()\n        \n        batch_size = inputs['input_ids'].shape[0]\n        batch_notebook_cell_ids = test_df.iloc[notebook_idx:notebook_idx+batch_size].merged_cell_ids.values\n        for i in range(batch_size): \n            notebook_token_type_ids = inputs['token_type_ids'][i]\n            notebook_cell_ids = test_df.iloc[notebook_idx+i].merged_cell_ids.split(CELL_SEP)\n            notebook_token_preds = token_preds[i]\n\n            cur_cell_idx = 0\n            prev_token_type_id = 0\n            for token_pred, token_type_id in zip(notebook_token_preds, notebook_token_type_ids):\n                if token_type_id == -100:\n                    continue\n                if token_type_id == prev_token_type_id:\n                    cur_cell_id = notebook_cell_ids[cur_cell_idx]\n                    cell_id_to_rank_preds[cur_cell_id].append(token_pred)\n                else: \n                    cur_cell_idx += 1\n                prev_token_type_id = token_type_id\n        notebook_idx += batch_size\ncell_id_to_pred_rank = {cell_id: float(sum(preds)/len(preds)) for cell_id, preds in cell_id_to_rank_preds.items()}","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:50:05.056133Z","iopub.status.busy":"2022-07-02T01:50:05.055727Z","iopub.status.idle":"2022-07-02T01:50:07.203674Z","shell.execute_reply":"2022-07-02T01:50:07.202765Z"},"papermill":{"duration":2.158067,"end_time":"2022-07-02T01:50:07.205605","exception":false,"start_time":"2022-07-02T01:50:05.047538","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample_sub = pd.read_csv('/kaggle/input/AI4Code/sample_submission.csv')\npredicted_cell_orders = []\nfor notebook_idx in range(len(test_df)):\n    cell_ids = test_df.iloc[notebook_idx].merged_cell_ids.split(CELL_SEP)\n    pred_cell_ranks = [\n        cell_id_to_pred_rank.get(cell_id, cell_idx/len(cell_ids)) \n        for cell_idx, cell_id in enumerate(cell_ids)\n    ]\n    ordered_cell_ids_list = [cell_id for cell_pred, cell_id in sorted(zip(pred_cell_ranks, cell_ids), key=lambda pairs: pairs[0])]\n    ordered_cell_ids = ' '.join(ordered_cell_ids_list)\n    predicted_cell_orders.append(ordered_cell_ids)\ntest_df['cell_order'] = predicted_cell_orders\nif HP.hide_lb_score: \n    test_df.cell_order = test_df.cell_order.apply(lambda cell_order: ' '.join(cell_order.split()[::-1]))\ntest_df['id'] = test_df.notebook_id\ntest_df[['id', 'cell_order']].to_csv('submission.csv', index=False)","metadata":{"execution":{"iopub.execute_input":"2022-07-02T01:50:09.223646Z","iopub.status.busy":"2022-07-02T01:50:09.223039Z","iopub.status.idle":"2022-07-02T01:50:09.350915Z","shell.execute_reply":"2022-07-02T01:50:09.349896Z"},"papermill":{"duration":2.140113,"end_time":"2022-07-02T01:50:09.353257","exception":false,"start_time":"2022-07-02T01:50:07.213144","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"papermill":{"duration":0.012379,"end_time":"2022-07-02T01:50:09.381729","exception":false,"start_time":"2022-07-02T01:50:09.36935","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]}]}