{"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: Longformer 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 Loop </a> \n","metadata":{"papermill":{"duration":0.008982,"end_time":"2022-07-21T04:25:58.140348","exception":false,"start_time":"2022-07-21T04:25:58.131366","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# TODO: Torch Profiler\n# todo: Better batch selection and offline tokenixation\n# TODO: MAKE SURE THAT TOKEN CELL INDICES IS POSITIVE\n\n# 0.674 empty w. 4096","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:25:58.157123Z","iopub.status.busy":"2022-07-21T04:25:58.156677Z","iopub.status.idle":"2022-07-21T04:25:58.162165Z","shell.execute_reply":"2022-07-21T04:25:58.161262Z"},"papermill":{"duration":0.016612,"end_time":"2022-07-21T04:25:58.164502","exception":false,"start_time":"2022-07-21T04:25:58.147890","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sync Notebook with VS Code #\nimport sys\nsys.path.append('/kaggle/input/github-ai4code/ai4code')\nsys.path.append('/kaggle/input/omegaconf')\n!cp -r /kaggle/input/github-ai4code/ai4code /kaggle/working\n\n# Run Setup Scripts #\n%run /kaggle/working/ai4code/ai4c/jupyter_setup.py\n\n# Imports #\nimport ai4c\nimport ai4c.process_df\nimport torch","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:25:58.180719Z","iopub.status.busy":"2022-07-21T04:25:58.179907Z","iopub.status.idle":"2022-07-21T04:26:04.796660Z","shell.execute_reply":"2022-07-21T04:26:04.795585Z"},"papermill":{"duration":6.627344,"end_time":"2022-07-21T04:26:04.799148","exception":false,"start_time":"2022-07-21T04:25:58.171804","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 Loop </a> \n\n<a name='hyperparameters'>","metadata":{"papermill":{"duration":0.007321,"end_time":"2022-07-21T04:26:04.814500","exception":false,"start_time":"2022-07-21T04:26:04.807179","status":"completed"},"tags":[]}},{"cell_type":"code","source":"%%hyperparameters HP\n\n## Huggingface Backbone ##\nbackbone_name: 'allenai/longformer-base-4096'\nbackbone_folder: 'longformer-backbones'\n\nattention_probs_dropout_prob: 0.10\nhidden_dropout_prob: 0.10\n\ngradient_checkpointing: True\n\n\n## Tokenization & Pre-processing ##\nmax_seq_len: 1024\nmax_markdown_seq_len: 512\nmax_tokens_per_cell: 512\nmax_global_tokens_per_notebook: 128\n\n\n## Data Factory ##\ntrain_folds: [1]\nvalid_fold: 0\nnum_validation_notebooks: 1000\nsort_notebooks_by_input_tokens: True\n\n\n## Model Training ##\nnum_train_epochs: 1\ntrain_batch_size: 4\neval_batch_size: 4\n\ngradient_accumulation_steps: 16\nmixed_precision: True\n\n\n## Loss Function ##\nmarkdown_cell_loss_weight: 0.50\nloss_fn_name: 'mse'\n\n\n## Cosine Decay LR Scheduler ##\nwarmup_ratio: 0.0625\nlearning_rate: 3e-5\n\n\n## AdamW Optimizer ##\nweight_decay: 1e-4\nmax_grad_norm: 1e6\nadam_epsilon: 1e-6\n\n\n## Load From Cache: Tokenized Dataset ##\nprocessed_dataset_folder: null # 'ai4code-flax-seq2seq-tokenization-2048'\ndebug_notebooks: 10000\n\n\n## Logging ##\nlogging_frequency: 10\n\n## Global Args ##\nhide_lb_score: False","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:04.831532Z","iopub.status.busy":"2022-07-21T04:26:04.830844Z","iopub.status.idle":"2022-07-21T04:26:04.885114Z","shell.execute_reply":"2022-07-21T04:26:04.884281Z"},"papermill":{"duration":0.065357,"end_time":"2022-07-21T04:26:04.887210","exception":false,"start_time":"2022-07-21T04:26:04.821853","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"backbone_code = HP.backbone_name.replace('/', '_')\nbackbone_dir = f'/kaggle/input/{HP.backbone_folder}/{backbone_code}'\nprint(f'Loading backbone and tokenizer from {backbone_dir}')\n\nbackbone = transformers.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    gradient_checkpointing=HP.gradient_checkpointing,\n)\ntokenizer = transformers.AutoTokenizer.from_pretrained(backbone_dir, use_fast=True)","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:04.903267Z","iopub.status.busy":"2022-07-21T04:26:04.902996Z","iopub.status.idle":"2022-07-21T04:26:12.446509Z","shell.execute_reply":"2022-07-21T04:26:12.442060Z"},"papermill":{"duration":7.554993,"end_time":"2022-07-21T04:26:12.449759","exception":false,"start_time":"2022-07-21T04:26:04.894766","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## ⚒️ Data Factory ⚒️\n\n---\n#### <a href='#prepare-huggingface-datasets'> 🤗 Huggingface Datasets </a> | <a href='#prepare-pytorch-datasets'> 🔥 PyTorch Data Module </a> \n\n\n<a name='data-factory'>","metadata":{"papermill":{"duration":0.027101,"end_time":"2022-07-21T04:26:12.499507","exception":false,"start_time":"2022-07-21T04:26:12.472406","status":"completed"},"tags":[]}},{"cell_type":"code","source":"processed_dataset_path = Path(f'/kaggle/input/{HP.processed_dataset_folder}')\nif HP.processed_dataset_folder is not None:\n    print(f'Loading dataframes from {processed_dataset_path}')\n    notebooks_df = pd.read_csv('/kaggle/input/ai4code-dataframes/notebooks_df.csv')\n    train_df = notebooks_df[notebooks_df.notebook_fold != HP.valid_fold]\n    valid_df = notebooks_df[notebooks_df.notebook_fold == HP.valid_fold]\n    valid_df = valid_df.sample(HP.num_validation_notebooks)\nelse:\n    print(f'Loading {HP.debug_notebooks} notebooks for debugging.')\n    train_df = valid_df = notebooks_df = pd.read_csv('/kaggle/input/ai4code-dataframes/notebooks_df.csv', nrows=HP.debug_notebooks)","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:12.537033Z","iopub.status.busy":"2022-07-21T04:26:12.536550Z","iopub.status.idle":"2022-07-21T04:26:18.625486Z","shell.execute_reply":"2022-07-21T04:26:18.624499Z"},"papermill":{"duration":6.114361,"end_time":"2022-07-21T04:26:18.628198","exception":false,"start_time":"2022-07-21T04:26:12.513837","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Fast Submission #\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)\n\nif len(test_df) < 100:\n    train_df = valid_df = train_df.sample(64)","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:18.646135Z","iopub.status.busy":"2022-07-21T04:26:18.645851Z","iopub.status.idle":"2022-07-21T04:26:18.853955Z","shell.execute_reply":"2022-07-21T04:26:18.852916Z"},"papermill":{"duration":0.220115,"end_time":"2022-07-21T04:26:18.856703","exception":false,"start_time":"2022-07-21T04:26:18.636588","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 🤗 Prepare Huggingface Datasets\n---\n#### <a href='#data-factory'> ⚒ Data Factory </a>  | <a href='#hyperparameters'> ⚙️ Hyperparameters </a>|  <a href='#training'> ⚡ Training </a> \n\n<a name='prepare-huggingface-datasets'>","metadata":{"papermill":{"duration":0.007577,"end_time":"2022-07-21T04:26:18.872520","exception":false,"start_time":"2022-07-21T04:26:18.864943","status":"completed"},"tags":[]}},{"cell_type":"code","source":"%%writefile prepare_hf_dataset.py\n\nfrom functools import partial\nfrom tqdm.auto import tqdm\nimport pandas as pd\nimport numpy as np\nimport argparse\n\nimport transformers\nimport datasets\n\ntqdm.pandas()\nCELL_SEP = '[CELL_SEP]'\n\ndef prune_cell_tokens(cell_token_ids, max_seq_len):\n    \"\"\"\n    Prunes cells that take too many tokens to fit in max_seq_len.\n    \"\"\"\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_longformer(\n    notebook_dict,\n    tokenizer,\n    max_seq_len,\n    max_markdown_seq_len,\n    max_tokens_per_cell,\n    max_global_tokens_per_notebook,\n):\n    '''Tokenize the notebook and convert to features for the model'''\n\n    markdown_cell_sources = notebook_dict['merged_markdown_cell_sources'].split(CELL_SEP)\n    markdown_cell_pct_ranks = [float(rank) for rank in notebook_dict['merged_markdown_cell_pct_ranks'].split(CELL_SEP)]\n    markdown_cell_ids = notebook_dict['merged_markdown_cell_ids'].split(CELL_SEP)\n\n    code_cell_sources = notebook_dict['merged_code_cell_sources'].split(CELL_SEP)\n    code_cell_pct_ranks = [float(rank) for rank in notebook_dict['merged_code_cell_pct_ranks'].split(CELL_SEP)]\n    code_cell_ids = notebook_dict['merged_code_cell_ids'].split(CELL_SEP)\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        markdown_cell_ids = markdown_cell_ids[: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        code_cell_ids = code_cell_ids[:max_code_cells]\n    \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    cell_pct_ranks = markdown_cell_pct_ranks + code_cell_pct_ranks\n    \n    input_ids, markdown_token_mask, code_token_mask = [], [], []\n    global_attention_mask = []\n    token_weights, token_labels = [], []\n    token_cell_indices = []\n    \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            markdown_token_mask += [1]*token_count_for_cell\n            code_token_mask += [0]*token_count_for_cell\n        else: \n            markdown_token_mask += [0]*token_count_for_cell\n            code_token_mask += [1]*token_count_for_cell\n        \n        if sum(global_attention_mask) < max_global_tokens_per_notebook:\n            global_attention_mask += [1] + [0]*(token_count_for_cell-1)\n        else: \n            global_attention_mask += [0]*token_count_for_cell\n        input_ids += cell_token_ids\n        token_cell_indices += [cur_cell_idx] * token_count_for_cell\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    # Pad to max_seq_len for efficient storage \n    num_pad_tokens = max_seq_len - len(input_ids)\n    attention_mask = [1]*len(input_ids) + [0]*num_pad_tokens\n    input_ids += [0]*num_pad_tokens\n    global_attention_mask += [0]*num_pad_tokens\n    markdown_token_mask += [0]*num_pad_tokens\n    code_token_mask += [0]*num_pad_tokens\n    token_cell_indices += [-100]*num_pad_tokens\n    token_labels += [-100]*num_pad_tokens\n    token_weights += [0]*num_pad_tokens\n\n    # Build the feature dict for the input \n    notebook_features = {\n        'input_ids': input_ids, \n        'attention_mask': attention_mask,\n        'global_attention_mask': global_attention_mask,\n        'markdown_token_mask': markdown_token_mask,\n        'code_token_mask': code_token_mask,\n        'token_cell_indices': token_cell_indices,\n        \n        'token_labels': token_labels,\n        'token_weights': token_weights,\n        \n        'notebook_id': notebook_dict['notebook_id'],\n        'num_pad_tokens': num_pad_tokens,\n    }\n    return notebook_features\n\n\ndef build_hf_dataset(\n    df, \n    tokenizer, \n    max_seq_len,\n    max_markdown_seq_len,\n    max_tokens_per_cell,\n    max_global_tokens_per_notebook,\n    ):\n    '''Builds the huggingface dataset for training the model.'''\n    convert_to_features = partial(\n        convert_to_features_longformer, \n        tokenizer=tokenizer,\n        max_seq_len=max_seq_len,\n        max_markdown_seq_len=max_markdown_seq_len,\n        max_tokens_per_cell=max_tokens_per_cell,\n        max_global_tokens_per_notebook=max_global_tokens_per_notebook,\n    )\n    raw_dataset = datasets.Dataset.from_pandas(df)\n    processed_dataset = raw_dataset.map(\n        convert_to_features, \n        remove_columns=raw_dataset.column_names, \n        desc='Running tokenizer on raw dataset'\n    )\n    processed_dataset.set_format(type='numpy')\n    empty_sentences = (np.array(processed_dataset['attention_mask'])[:, -1] == 0).sum()\n    print('Empty sentences ratio:', empty_sentences/len(processed_dataset))\n    return processed_dataset\n\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--tokenizer_name', default='google/bigbird-roberta-large', type=str, help='The tokenizer name')\n    parser.add_argument('--max_seq_len', default=512, type=int, help='The maximum sequence length')\n    parser.add_argument('--max_markdown_seq_len', default=512, type=int, help='The maximum sequence length for markdown cells')\n    parser.add_argument('--max_tokens_per_cell', default=256, type=int, help='The maximum number of tokens per cell')\n    parser.add_argument('--max_global_tokens_per_notebook', default=128, type=int, help='The maximum number of global tokens per notebook')\n    parser.add_argument('--notebooks_df_path', default='notebooks_df.csv', type=str, help='Path to notebooks.csv')\n\n    args = parser.parse_args()\n    \n    tokenizer = transformers.AutoTokenizer.from_pretrained(args.tokenizer_name)\n    notebooks_df = pd.read_csv(args.notebooks_df_path)\n    print('Total number of notebooks:', len(notebooks_df))\n    \n    for fold in tqdm(range(8), desc='Tokenizing notebooks for each fold'):\n        fold_df = notebooks_df[notebooks_df.notebook_fold == fold]\n        fold_dataset = build_hf_dataset(\n            df=fold_df,\n            tokenizer=tokenizer,\n            max_seq_len=args.max_seq_len,\n            max_markdown_seq_len=args.max_markdown_seq_len,\n            max_tokens_per_cell=args.max_tokens_per_cell,\n            max_global_tokens_per_notebook=args.max_global_tokens_per_notebook,\n        )\n        fold_dataset.save_to_disk(f'hf_dataset_fold_{fold}')\n    print('Done!')","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:18.890412Z","iopub.status.busy":"2022-07-21T04:26:18.890120Z","iopub.status.idle":"2022-07-21T04:26:18.944467Z","shell.execute_reply":"2022-07-21T04:26:18.943222Z"},"papermill":{"duration":0.065827,"end_time":"2022-07-21T04:26:18.946448","exception":false,"start_time":"2022-07-21T04:26:18.880621","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import prepare_hf_dataset\n\nif HP.processed_dataset_folder is not None:\n    print(f'Loading tokenized datasets from {processed_dataset_path}')\n    valid_hf_dataset = datasets.load_from_disk(processed_dataset_path/f'hf_dataset_fold_{HP.valid_fold}')\n    train_hf_dataset = datasets.concatenate_datasets([\n        datasets.load_from_disk(processed_dataset_path/f'hf_dataset_fold_{fold}')\n        for fold in tqdm(range(8), desc='Loading training dataset')\n        if fold != HP.valid_fold\n    ])\n    train_hf_dataset.save_to_disk('train_hf_dataset')\n    train_hf_dataset = datasets.load_from_disk('train_hf_dataset')\nelse:\n    train_hf_dataset = valid_hf_dataset = prepare_hf_dataset.build_hf_dataset(\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_global_tokens_per_notebook=HP.max_global_tokens_per_notebook,\n        max_tokens_per_cell=HP.max_tokens_per_cell,\n    )\n\nif HP.sort_notebooks_by_input_tokens:\n    train_hf_dataset = train_hf_dataset.sort('num_pad_tokens', reverse=True)\n    valid_hf_dataset = valid_hf_dataset.sort('num_pad_tokens', reverse=True)","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:18.963528Z","iopub.status.busy":"2022-07-21T04:26:18.963213Z","iopub.status.idle":"2022-07-21T04:26:20.383238Z","shell.execute_reply":"2022-07-21T04:26:20.382324Z"},"papermill":{"duration":1.431024,"end_time":"2022-07-21T04:26:20.385413","exception":false,"start_time":"2022-07-21T04:26:18.954389","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 🔥 Prepare PyTorch Data Module\n---\n#### <a href='#data-factory'> ⚒ Data Factory </a>  | <a href='#hyperparameters'> ⚙️ Hyperparameters </a>|  <a href='#training'> ⚡ Training </a> \n\n<a name='prepare-tensorflow-datasets'>","metadata":{"papermill":{"duration":0.007917,"end_time":"2022-07-21T04:26:20.401718","exception":false,"start_time":"2022-07-21T04:26:20.393801","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def collate_fn(batch):\n    min_pad_tokens_in_batch = min(inputs['num_pad_tokens'] for inputs in batch)\n    max_seq_len = len(batch[0]['input_ids'])\n    batch_seq_len = max_seq_len - min_pad_tokens_in_batch\n\n    # Prune the elements in the batch to batch_seq_len\n    model_inputs_batch = {\n        'input_ids': torch.tensor([inputs['input_ids'][:batch_seq_len] for inputs in batch], dtype=torch.long),\n        'attention_mask': torch.tensor([inputs['attention_mask'][:batch_seq_len] for inputs in batch], dtype=torch.long),\n        'global_attention_mask': torch.tensor([inputs['global_attention_mask'][:batch_seq_len] for inputs in batch], dtype=torch.long),\n        'markdown_token_mask': torch.tensor([inputs['markdown_token_mask'][:batch_seq_len] for inputs in batch], dtype=torch.long),\n        'code_token_mask': torch.tensor([inputs['code_token_mask'][:batch_seq_len] for inputs in batch], dtype=torch.long),\n    }\n\n    model_labels_batch = {\n        'token_labels': torch.tensor([inputs['token_labels'][:batch_seq_len] for inputs in batch], dtype=torch.long),\n        'token_weights': torch.tensor([inputs['token_weights'][:batch_seq_len] for inputs in batch], dtype=torch.float),\n    }\n    return model_inputs_batch, model_labels_batch\n\nshuffle_train_dataset = not HP.sort_notebooks_by_input_tokens\ntrain_dataloader = torch.utils.data.DataLoader(\n    dataset=train_hf_dataset,\n    batch_size=HP.train_batch_size,\n    shuffle=shuffle_train_dataset,\n    num_workers=2,\n    collate_fn=collate_fn,\n    pin_memory=True,\n    drop_last=True,\n    prefetch_factor=4,\n)\n\neval_dataloader = torch.utils.data.DataLoader(\n    dataset=valid_hf_dataset,\n    batch_size=HP.eval_batch_size,\n    shuffle=False,\n    num_workers=2,\n    collate_fn=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-21T04:26:20.419433Z","iopub.status.busy":"2022-07-21T04:26:20.418631Z","iopub.status.idle":"2022-07-21T04:26:23.491818Z","shell.execute_reply":"2022-07-21T04:26:23.490695Z"},"papermill":{"duration":3.084578,"end_time":"2022-07-21T04:26:23.494257","exception":false,"start_time":"2022-07-21T04:26:20.409679","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, global_attention_mask):\n        backbone_outputs = self.backbone(\n            input_ids=input_ids,\n            global_attention_mask=global_attention_mask,\n        )\n        token_preds = self.ranker(backbone_outputs.last_hidden_state)\n        batch_seq_len = token_preds.size(1)\n        return token_preds.view((-1, batch_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-21T04:26:23.513699Z","iopub.status.busy":"2022-07-21T04:26:23.513097Z","iopub.status.idle":"2022-07-21T04:26:23.568920Z","shell.execute_reply":"2022-07-21T04:26:23.567528Z"},"papermill":{"duration":0.068305,"end_time":"2022-07-21T04:26:23.571452","exception":false,"start_time":"2022-07-21T04:26:23.503147","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch_model\nimport gc\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)\ntrain_steps_per_epoch = math.ceil(len(train_dataloader) / HP.gradient_accumulation_steps)\ntotal_train_steps = HP.num_train_epochs * train_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_ = model.cuda()\n\ntorch.cuda.empty_cache()\ngc.collect()","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:23.589850Z","iopub.status.busy":"2022-07-21T04:26:23.589573Z","iopub.status.idle":"2022-07-21T04:26:29.530257Z","shell.execute_reply":"2022-07-21T04:26:29.529317Z"},"papermill":{"duration":5.952256,"end_time":"2022-07-21T04:26:29.532540","exception":false,"start_time":"2022-07-21T04:26:23.580284","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# ⚡ Training Loop ⚡\n---\n### <a href='#hyperparameters'> ⚙️ Hyperparameters </a>  | <a href='#model'> 🧠 Model </a>\n\n<a name='training'>","metadata":{"papermill":{"duration":0.008373,"end_time":"2022-07-21T04:26:29.549846","exception":false,"start_time":"2022-07-21T04:26:29.541473","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def mse_loss(token_labels, token_preds, token_weights, token_mask):\n    token_labels, token_preds, token_weights = token_labels*token_mask, token_preds*token_mask, token_weights*token_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\n    return torch.mean(notebook_losses)\n\n_ = model.train()\nscaler = torch.cuda.amp.GradScaler(enabled=HP.mixed_precision)\n\nrunning_metrics = defaultdict(int)\nsteps_progress_bar = tqdm(range(total_train_steps), desc=\"Training Progress\") \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            global_attention_mask=inputs['global_attention_mask'],\n        )\n        \n        markdown_cell_loss = mse_loss(\n            token_labels=labels['token_labels'],\n            token_preds=token_preds,\n            token_weights=labels['token_weights'],\n            token_mask=inputs['markdown_token_mask'],\n        )\n        code_cell_loss = mse_loss(\n            token_labels=labels['token_labels'],\n            token_preds=token_preds,\n            token_weights=labels['token_weights'],\n            token_mask=inputs['code_token_mask'],\n        )\n        loss = HP.markdown_cell_loss_weight*markdown_cell_loss + (1-HP.markdown_cell_loss_weight) * code_cell_loss\n    \n    grad_norm = torch.nn.utils.clip_grad_norm_(model.backbone.parameters(), HP.max_grad_norm)\n    \n    running_metrics['total_loss'] += loss.item()\n    running_metrics['markdown_cell_loss'] += markdown_cell_loss.item()\n    running_metrics['code_cell_loss'] += code_cell_loss.item()\n    running_metrics['learning_rate'] += lr_scheduler.get_lr()[0]\n    running_metrics['num_input_tokens'] += torch.sum(inputs['attention_mask']).item()/HP.train_batch_size\n    running_metrics['gradient_norm'] += grad_norm.item()\n\n    steps_progress_bar.set_postfix(\n        loss=f\"{loss.item()/HP.gradient_accumulation_steps:.4f}\", \n        markdown_cell_loss=f\"{markdown_cell_loss.item()/HP.gradient_accumulation_steps:.4f}\", \n        code_cell_loss=f\"{code_cell_loss.item()/HP.gradient_accumulation_steps:.4f}\",\n        grad_norm=f\"{grad_norm.item():.2f}\",\n    )\n\n    scaler.scale(loss).backward()\n    if (step+1) % HP.gradient_accumulation_steps == 0:\n        scaler.step(optimizer)\n        scaler.update()\n        lr_scheduler.step()\n        optimizer.zero_grad()\n        _ = steps_progress_bar.update(1)\n    \n    logging_width = (HP.logging_frequency * HP.gradient_accumulation_steps)\n    if (step+1) % logging_width == 0:\n        print('-'*50)\n        print(f\"Step {step-logging_width}-{step} out of {total_train_steps}\")\n        for k, v in running_metrics.items():\n            print(colored(k, 'blue'), ':', colored(v / logging_width, 'red'))\n        running_metrics = defaultdict(int)\n        print()","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:29.569459Z","iopub.status.busy":"2022-07-21T04:26:29.568061Z","iopub.status.idle":"2022-07-21T04:26:49.319340Z","shell.execute_reply":"2022-07-21T04:26:49.318243Z"},"papermill":{"duration":19.763376,"end_time":"2022-07-21T04:26:49.321693","exception":false,"start_time":"2022-07-21T04:26:29.558317","status":"completed"},"tags":[],"collapsed":true,"jupyter":{"outputs_hidden":true}},"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.009186,"end_time":"2022-07-21T04:26:49.340923","exception":false,"start_time":"2022-07-21T04:26:49.331737","status":"completed"},"tags":[]}},{"cell_type":"code","source":"from collections import defaultdict\n\nCELL_SEP = '[CELL_SEP]'\ntorch.cuda.empty_cache()\ngc.collect()\n_ = model.eval()\n\nepoch_metrics = defaultdict(int)\ncell_id_to_pred_rank = {}\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                global_attention_mask=inputs['global_attention_mask'],\n            )\n            markdown_cell_loss = mse_loss(\n                token_labels=labels['token_labels'],\n                token_preds=token_preds,\n                token_weights=labels['token_weights'],\n                token_mask=inputs['markdown_token_mask'],\n            )\n            code_cell_loss = mse_loss(\n                token_labels=labels['token_labels'],\n                token_preds=token_preds,\n                token_weights=labels['token_weights'],\n                token_mask=inputs['code_token_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        notebook_ids = valid_hf_dataset['notebook_id'][step*batch_size: (step+1)*batch_size]\n        token_cell_indices = valid_hf_dataset['token_cell_indices'][step*batch_size: (step+1)*batch_size]\n        for example_idx in range(batch_size):\n            # POTENTIAL BUG\n            \n            notebook_id = notebook_ids[example_idx]\n            notebook_row = valid_df[valid_df.notebook_id==notebook_id].iloc[0]\n\n            cell_ids = notebook_row.merged_cell_ids.split(CELL_SEP)\n            example_token_preds = token_preds[example_idx]\n            example_token_cell_indices = token_cell_indices[example_idx]\n\n            cell_idx_to_sum_preds = defaultdict(int)\n            cell_idx_to_num_preds = defaultdict(int)\n            for token_pred, cell_idx in zip(example_token_preds, example_token_cell_indices):\n                if cell_idx < 0:\n                    break\n                cell_idx_to_sum_preds[cell_idx] += token_pred\n                cell_idx_to_num_preds[cell_idx] += 1\n            \n            for cell_idx, sum_pred in cell_idx_to_sum_preds.items():\n                num_pred = cell_idx_to_num_preds[cell_idx]\n                cell_id_to_pred_rank[cell_ids[cell_idx]] = sum_pred / num_pred\n                \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\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    cutoff_df = valid_df[valid_df.markdown_cell_count>cutoff]\n    tau = cutoff_df.kendall_tau.mean()\n    print(f\"Kendall Tau for {colored(len(cutoff_df), 'blue')} notebooks with {colored(cutoff, 'yellow')}+ markdown cells:\", \\\n          colored(tau, 'red'))\n\nvalid_df.to_csv(f'valid_df.csv', index=False)","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:49.361975Z","iopub.status.busy":"2022-07-21T04:26:49.361015Z","iopub.status.idle":"2022-07-21T04:26:54.680333Z","shell.execute_reply":"2022-07-21T04:26:54.679227Z"},"papermill":{"duration":5.332977,"end_time":"2022-07-21T04:26:54.683042","exception":false,"start_time":"2022-07-21T04:26:49.350065","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# tau_32 = valid_df[valid_df.markdown_cell_count>32].kendall_tau.mean()\n# model_save_path = f\"{backbone_code}_tau{int(avg_tau*10000)}_tau32_{int(tau_32*1000)}.pt\"\n# torch.save(model.state_dict(), model_save_path)\n# print(f'Model saved at {model_save_path}')","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:54.704875Z","iopub.status.busy":"2022-07-21T04:26:54.704181Z","iopub.status.idle":"2022-07-21T04:26:54.775459Z","shell.execute_reply":"2022-07-21T04:26:54.774519Z"},"papermill":{"duration":0.083875,"end_time":"2022-07-21T04:26:54.777337","exception":false,"start_time":"2022-07-21T04:26:54.693462","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.009557,"end_time":"2022-07-21T04:26:54.796388","exception":false,"start_time":"2022-07-21T04:26:54.786831","status":"completed"},"tags":[]}},{"cell_type":"code","source":"del train_df, valid_df, train_hf_dataset, valid_hf_dataset\ntorch.cuda.empty_cache()\n\ngc.collect()","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:54.816581Z","iopub.status.busy":"2022-07-21T04:26:54.816152Z","iopub.status.idle":"2022-07-21T04:26:55.133307Z","shell.execute_reply":"2022-07-21T04:26:55.132396Z"},"papermill":{"duration":0.329853,"end_time":"2022-07-21T04:26:55.135683","exception":false,"start_time":"2022-07-21T04:26:54.805830","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if 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)","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:55.158056Z","iopub.status.busy":"2022-07-21T04:26:55.156551Z","iopub.status.idle":"2022-07-21T04:26:55.332122Z","shell.execute_reply":"2022-07-21T04:26:55.331087Z"},"papermill":{"duration":0.188633,"end_time":"2022-07-21T04:26:55.334267","exception":false,"start_time":"2022-07-21T04:26:55.145634","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%writefile prepare_hf_dataset.py\n\nfrom functools import partial\nfrom tqdm.auto import tqdm\nimport pandas as pd\nimport numpy as np\nimport argparse\n\nimport transformers\nimport datasets\n\ntqdm.pandas()\nCELL_SEP = '[CELL_SEP]'\n\ndef prune_cell_tokens(cell_token_ids, max_seq_len):\n    \"\"\"\n    Prunes cells that take too many tokens to fit in max_seq_len.\n    \"\"\"\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_longformer(\n    notebook_dict,\n    tokenizer,\n    max_seq_len,\n    max_markdown_seq_len,\n    max_tokens_per_cell,\n    max_global_tokens_per_notebook,\n):\n    '''Tokenize the notebook and convert to features for the model'''\n\n    markdown_cell_sources = notebook_dict['merged_markdown_cell_sources'].split(CELL_SEP)\n    markdown_cell_ids = notebook_dict['merged_markdown_cell_ids'].split(CELL_SEP)\n\n    code_cell_sources = notebook_dict['merged_code_cell_sources'].split(CELL_SEP)\n    code_cell_ids = notebook_dict['merged_code_cell_ids'].split(CELL_SEP)\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_ids = markdown_cell_ids[: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_ids = code_cell_ids[:max_code_cells]\n    \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    \n    input_ids, markdown_token_mask, code_token_mask = [], [], []\n    global_attention_mask = []\n    token_cell_indices = []\n    \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            markdown_token_mask += [1]*token_count_for_cell\n            code_token_mask += [0]*token_count_for_cell\n        else: \n            markdown_token_mask += [0]*token_count_for_cell\n            code_token_mask += [1]*token_count_for_cell\n        \n        if sum(global_attention_mask) < max_global_tokens_per_notebook:\n            global_attention_mask += [1] + [0]*(token_count_for_cell-1)\n        else: \n            global_attention_mask += [0]*token_count_for_cell\n        input_ids += cell_token_ids\n        token_cell_indices += [cur_cell_idx] * token_count_for_cell\n    \n    # Pad to max_seq_len for efficient storage \n    num_pad_tokens = max_seq_len - len(input_ids)\n    attention_mask = [1]*len(input_ids) + [0]*num_pad_tokens\n    input_ids += [0]*num_pad_tokens\n    global_attention_mask += [0]*num_pad_tokens\n    markdown_token_mask += [0]*num_pad_tokens\n    code_token_mask += [0]*num_pad_tokens\n    token_cell_indices += [-100]*num_pad_tokens\n\n    # Build the feature dict for the input \n    notebook_features = {\n        'input_ids': input_ids, \n        'attention_mask': attention_mask,\n        'global_attention_mask': global_attention_mask,\n        'markdown_token_mask': markdown_token_mask,\n        'code_token_mask': code_token_mask,\n        'token_cell_indices': token_cell_indices,\n        \n        'notebook_id': notebook_dict['notebook_id'],\n        'num_pad_tokens': num_pad_tokens,\n    }\n    return notebook_features\n\n\ndef build_hf_dataset(\n    df, \n    tokenizer, \n    max_seq_len,\n    max_markdown_seq_len,\n    max_tokens_per_cell,\n    max_global_tokens_per_notebook,\n    ):\n    '''Builds the huggingface dataset for training the model.'''\n    convert_to_features = partial(\n        convert_to_features_longformer, \n        tokenizer=tokenizer,\n        max_seq_len=max_seq_len,\n        max_markdown_seq_len=max_markdown_seq_len,\n        max_tokens_per_cell=max_tokens_per_cell,\n        max_global_tokens_per_notebook=max_global_tokens_per_notebook,\n    )\n    raw_dataset = datasets.Dataset.from_pandas(df)\n    processed_dataset = raw_dataset.map(\n        convert_to_features, \n        remove_columns=raw_dataset.column_names, \n        desc='Running tokenizer on raw dataset'\n    )\n    processed_dataset.set_format(type='numpy')\n    empty_sentences = (np.array(processed_dataset['attention_mask'])[:, -1] == 0).sum()\n    print('Empty sentences ratio:', empty_sentences/len(processed_dataset))\n    return processed_dataset","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:55.356728Z","iopub.status.busy":"2022-07-21T04:26:55.356463Z","iopub.status.idle":"2022-07-21T04:26:55.422776Z","shell.execute_reply":"2022-07-21T04:26:55.421583Z"},"papermill":{"duration":0.07967,"end_time":"2022-07-21T04:26:55.424891","exception":false,"start_time":"2022-07-21T04:26:55.345221","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_hf_dataset = prepare_hf_dataset.build_hf_dataset(\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_global_tokens_per_notebook=HP.max_global_tokens_per_notebook,\n    max_tokens_per_cell=HP.max_tokens_per_cell,\n)","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:55.445826Z","iopub.status.busy":"2022-07-21T04:26:55.445571Z","iopub.status.idle":"2022-07-21T04:26:55.606150Z","shell.execute_reply":"2022-07-21T04:26:55.605141Z"},"papermill":{"duration":0.173278,"end_time":"2022-07-21T04:26:55.608112","exception":false,"start_time":"2022-07-21T04:26:55.434834","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def collate_fn(batch):\n    min_pad_tokens_in_batch = min(inputs['num_pad_tokens'] for inputs in batch)\n    max_seq_len = len(batch[0]['input_ids'])\n    batch_seq_len = max_seq_len - min_pad_tokens_in_batch\n\n    # Prune the elements in the batch to batch_seq_len\n    model_inputs_batch = {\n        'input_ids': torch.tensor([inputs['input_ids'][:batch_seq_len] for inputs in batch], dtype=torch.long),\n        'attention_mask': torch.tensor([inputs['attention_mask'][:batch_seq_len] for inputs in batch], dtype=torch.long),\n        'global_attention_mask': torch.tensor([inputs['global_attention_mask'][:batch_seq_len] for inputs in batch], dtype=torch.long),\n        'markdown_token_mask': torch.tensor([inputs['markdown_token_mask'][:batch_seq_len] for inputs in batch], dtype=torch.long),\n        'code_token_mask': torch.tensor([inputs['code_token_mask'][:batch_seq_len] for inputs in batch], dtype=torch.long),\n    }\n    return model_inputs_batch\n\ntest_dataloader = torch.utils.data.DataLoader(\n    dataset=test_hf_dataset,\n    batch_size=HP.eval_batch_size,\n    shuffle=False,\n    num_workers=2,\n    collate_fn=collate_fn,\n    pin_memory=True,\n    drop_last=False,\n    prefetch_factor=4,\n)\n\nbatch = next(iter(test_dataloader))","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:55.632271Z","iopub.status.busy":"2022-07-21T04:26:55.630697Z","iopub.status.idle":"2022-07-21T04:26:55.908027Z","shell.execute_reply":"2022-07-21T04:26:55.906611Z"},"papermill":{"duration":0.292136,"end_time":"2022-07-21T04:26:55.911297","exception":false,"start_time":"2022-07-21T04:26:55.619161","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"_ = model.eval()\nCELL_SEP = '[CELL_SEP]'\n\ncell_id_to_pred_rank = {}\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                global_attention_mask=inputs['global_attention_mask'],\n            )\n            token_preds = token_preds.detach().cpu().numpy()\n        \n        batch_size = inputs['input_ids'].shape[0]\n        start = step*batch_size\n        end = min(start+batch_size, len(test_hf_dataset))\n        \n        notebook_ids = test_hf_dataset['notebook_id'][start: end]\n        token_cell_indices = test_hf_dataset['token_cell_indices'][start: end]\n        for example_idx, notebook_id in enumerate(notebook_ids):\n            notebook_row = test_df[test_df.notebook_id==notebook_id].iloc[0]\n\n            cell_ids = notebook_row.merged_cell_ids.split(CELL_SEP)\n            example_token_preds = token_preds[example_idx]\n            example_token_cell_indices = token_cell_indices[example_idx]\n\n            cell_idx_to_sum_preds = defaultdict(int)\n            cell_idx_to_num_preds = defaultdict(int)\n            for token_pred, cell_idx in zip(example_token_preds, example_token_cell_indices):\n                if cell_idx < 0:\n                    break\n                cell_idx_to_sum_preds[cell_idx] += token_pred\n                cell_idx_to_num_preds[cell_idx] += 1\n            \n            for cell_idx, sum_pred in cell_idx_to_sum_preds.items():\n                num_pred = cell_idx_to_num_preds[cell_idx]\n                cell_id_to_pred_rank[cell_ids[cell_idx]] = sum_pred / num_pred","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:55.950946Z","iopub.status.busy":"2022-07-21T04:26:55.949903Z","iopub.status.idle":"2022-07-21T04:26:56.581600Z","shell.execute_reply":"2022-07-21T04:26:56.579513Z"},"papermill":{"duration":0.653888,"end_time":"2022-07-21T04:26:56.583758","exception":false,"start_time":"2022-07-21T04:26:55.929870","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-21T04:26:56.607313Z","iopub.status.busy":"2022-07-21T04:26:56.607010Z","iopub.status.idle":"2022-07-21T04:26:56.690867Z","shell.execute_reply":"2022-07-21T04:26:56.689859Z"},"papermill":{"duration":0.097976,"end_time":"2022-07-21T04:26:56.692948","exception":false,"start_time":"2022-07-21T04:26:56.594972","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df","metadata":{"execution":{"iopub.execute_input":"2022-07-21T04:26:56.716952Z","iopub.status.busy":"2022-07-21T04:26:56.716097Z","iopub.status.idle":"2022-07-21T04:26:56.790708Z","shell.execute_reply":"2022-07-21T04:26:56.789708Z"},"papermill":{"duration":0.088884,"end_time":"2022-07-21T04:26:56.793123","exception":false,"start_time":"2022-07-21T04:26:56.704239","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"papermill":{"duration":0.010933,"end_time":"2022-07-21T04:26:56.816184","exception":false,"start_time":"2022-07-21T04:26:56.805251","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]}]}