{"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 Train: Big Bird 🐥__\n\n---\n### <a href='#hyperparameters'> ⚙️ Hyperparameters </a> | <a href='#data-factory'> ⚒ Data Factory </a>  | <a href='#training-loop'> ⚡ Training Loop </a>","metadata":{"papermill":{"duration":0.017049,"end_time":"2022-05-21T12:17:56.031848","exception":false,"start_time":"2022-05-21T12:17:56.014799","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# %%time\n# # TODO: Integrate WanDB\n# !pip install git+https://github.com/google-research/t5x\n# import t5x","metadata":{"execution":{"iopub.status.busy":"2022-07-21T09:15:12.217537Z","iopub.execute_input":"2022-07-21T09:15:12.218260Z","iopub.status.idle":"2022-07-21T09:17:13.374394Z","shell.execute_reply.started":"2022-07-21T09:15:12.218148Z","shell.execute_reply":"2022-07-21T09:17:13.373647Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# from flax.jax_utils import replicate, unreplicate\n# from flax import jax_utils, struct, traverse_util","metadata":{"execution":{"iopub.status.busy":"2022-07-14T21:51:17.774871Z","iopub.execute_input":"2022-07-14T21:51:17.775477Z","iopub.status.idle":"2022-07-14T21:51:17.780115Z","shell.execute_reply.started":"2022-07-14T21:51:17.775369Z","shell.execute_reply":"2022-07-14T21:51:17.779262Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%writefile flax_setup.py\n\nimport optax\nimport flax\nimport jax\n\nimport jax.numpy as jnp\nimport flax.linen as nn\n\nfrom flax.training import train_state\nfrom flax.training.common_utils import shard\n\nimport requests\nimport os\n\ndef kaggle_tpu_setup():\n    if 'TPU_NAME' not in os.environ:\n        print('TPU not found')\n        return\n    os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'\n    url = 'http:' + os.environ['TPU_NAME'].split(':')[1] + ':8475/requestversion/tpu_driver_nightly'\n    resp = requests.post(url)\n    jax.config.FLAGS.jax_xla_backend = 'tpu_driver'\n    jax.config.FLAGS.jax_backend_target = os.environ['TPU_NAME']\n    jax.config.update('jax_default_matmul_precision', 'bfloat16')\n\nkaggle_tpu_setup()","metadata":{"execution":{"iopub.status.busy":"2022-07-14T21:51:17.782043Z","iopub.execute_input":"2022-07-14T21:51:17.782303Z","iopub.status.idle":"2022-07-14T21:51:17.802832Z","shell.execute_reply.started":"2022-07-14T21:51:17.782274Z","shell.execute_reply":"2022-07-14T21:51:17.801523Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Installations for Flax & BigBird #\n!pip install --upgrade jaxlib jax flax optax -q\n!pip install --upgrade transformers -q\n\n# Sync Notebook with VS Code #\nimport sys; sys.path.append('ai4code')\n\nfrom ai4c.jupyter_setup import *\nfrom flax_setup import *\nimport ai4c","metadata":{"execution":{"iopub.status.busy":"2022-07-14T21:51:17.805032Z","iopub.execute_input":"2022-07-14T21:51:17.805700Z","iopub.status.idle":"2022-07-14T21:51:57.404977Z","shell.execute_reply.started":"2022-07-14T21:51:17.805650Z","shell.execute_reply":"2022-07-14T21:51:57.404102Z"},"trusted":true},"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-loop'> ⚡ Training Loop </a> \n\n<a name='hyperparameters'>","metadata":{"papermill":{"duration":0.018693,"end_time":"2022-05-21T12:18:38.858918","exception":false,"start_time":"2022-05-21T12:18:38.840225","status":"completed"},"tags":[]}},{"cell_type":"code","source":"%%hyperparameters HP\n\n## Huggingface Backbone ##\nbackbone_name: 'google/bigbird-roberta-base'\nbackbone_weights: null\n\nattention_probs_dropout_prob: 0.10\nhidden_dropout_prob: 0.10\n\nnum_random_blocks: 3\nblock_size: 64\n\ngradient_checkpointing: False\nmax_seq_length: 1536\n\n\n## Tokenization ##\nmax_markdown_seq_len: 1280\nmax_tokens_per_markdown_cell: 512\nmax_tokens_per_code_cell: 256\n\n\n## Model Training ##\nnum_train_epochs: 3\n\nper_device_train_batch_size: 2\nper_device_eval_batch_size: 2\n\n\n## Loss Function ##\nloss_fn_name: 'mse'\nmarkdown_cell_loss_weight: 0.50\n\n\n## Cosine Decay LR Scheduler ##\nwarmup_ratio: 0.10\npeak_lr: 1e-5\nmin_lr: 1e-8\n\n\n## AdamW Optimizer ## \nweight_decay: 1e-6\nmax_grad_norm: 1.00\nbeta_2: 0.98\nepsilon: 1e-6\nema_decay: 0.99\n\n\n# Load From Cache: Tokenized Dataset #\nprocessed_dataset_folder: 'ai4code-tokenization-bigbird'\ndebug_notebooks: null\n\n\n# Data Factory #\nvalid_fold: 0\nrandom_state: 69420\nlogging_freq: 100","metadata":{"execution":{"iopub.status.busy":"2022-07-14T21:51:57.406519Z","iopub.execute_input":"2022-07-14T21:51:57.406938Z","iopub.status.idle":"2022-07-14T21:51:57.475277Z","shell.execute_reply.started":"2022-07-14T21:51:57.406894Z","shell.execute_reply":"2022-07-14T21:51:57.474141Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## ⚒️ Data Factory ⚒️\n---\n#### <a href='#hyperparameters'> ⚙️ Hyperparameters </a>  | <a href='#model'> 🧠 Model </a>|  <a href='#training'> ⚡ Training </a> \n\n<a name='data-factory'>","metadata":{}},{"cell_type":"code","source":"tokenizer = transformers.AutoTokenizer.from_pretrained(HP.backbone_name)\n\nprocessed_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]\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.status.busy":"2022-07-14T21:51:57.477650Z","iopub.execute_input":"2022-07-14T21:51:57.477910Z","iopub.status.idle":"2022-07-14T21:53:35.915466Z","shell.execute_reply.started":"2022-07-14T21:51:57.477882Z","shell.execute_reply":"2022-07-14T21:53:35.914110Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 🤗 Huggingface Dataset\n---\n\n#### <a href='#hyperparameters'> ⚙️ Hyperparameters </a>  | <a href='#model'> 🧠 Model </a>|  <a href='#training-loop'> ⚡ Training Loop </a> ","metadata":{}},{"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\nimport gc\n\nimport transformers\nimport datasets\n\ntqdm.pandas()\nCELL_SEP = '[CELL_SEP]'\n\ndef prune_code_tokens(code_token_ids, max_seq_len):\n    \"\"\"\n    Prunes cells that take too many tokens to fit in max_seq_len.\n    \"\"\"\n    code_token_counts = [len(token_ids) for token_ids in code_token_ids]\n    total_number_of_cells = len(code_token_counts)\n    total_tokens_to_prune = max(sum(code_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_code_token_count = max(code_token_counts)\n        second_max_code_token_count = sorted(code_token_counts)[-2]\n        for cell_idx, code_token_count in enumerate(code_token_counts):\n            if not code_token_count == cur_max_code_token_count: \n                continue\n            \n            num_tokens_to_pop = min(code_token_count-second_max_code_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            code_token_counts[cell_idx] -= num_tokens_to_pop\n            break\n    \n    # Prune the cell tokens\n    pruned_code_token_ids = []\n    for code_token_ids, num_tokens_to_pop in zip(code_token_ids, tokens_to_prune_per_cell):\n        if num_tokens_to_pop == 0:\n            pruned_code_token_ids.append(code_token_ids)\n            continue\n        pruned_code_token_ids.append(code_token_ids[:-num_tokens_to_pop])\n    return pruned_code_token_ids\n\n\ndef convert_to_features_bigbird(\n    notebook_dict,\n    tokenizer,\n    max_seq_len,\n    max_markdown_seq_len,\n    max_tokens_per_markdown_cell,\n    max_tokens_per_code_cell,\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_markdown_cell, max_markdown_seq_len//markdown_cell_count)\n    markdown_code_token_ids = tokenizer(\n        markdown_cell_sources,\n        max_length=max_tokens_per_markdown_cell,\n        truncation=True,\n    )['input_ids']\n    markdown_code_token_ids = prune_code_tokens(markdown_code_token_ids, max_markdown_seq_len)\n    total_markdown_code_tokens = sum([len(token_ids) for token_ids in markdown_code_token_ids])\n\n    max_code_seq_len = max_seq_len - total_markdown_code_tokens\n    max_tokens_per_code_cell = max(max_tokens_per_code_cell, max_code_seq_len//code_cell_count)\n    code_code_token_ids = tokenizer(\n        code_cell_sources, \n        max_length=max_tokens_per_code_cell, \n        truncation=True, \n    )['input_ids']\n    code_code_token_ids = prune_code_tokens(code_code_token_ids, max_seq_len-total_markdown_code_tokens)\n\n    # Merge the tokenized cells and create the model features\n    code_token_ids = markdown_code_token_ids + code_code_token_ids\n    cell_ids = markdown_cell_ids + code_cell_ids\n    notebook_cell_count = len(code_token_ids)\n\n    # Create the model features\n    if 'merged_cell_pct_ranks' in notebook_dict:\n        cell_pct_ranks = markdown_cell_pct_ranks + code_cell_pct_ranks\n    else:\n        cell_pct_ranks = [-1]*notebook_cell_count\n    \n    input_ids, markdown_token_mask, code_token_mask = [], [], []\n    token_weights, token_labels = [], []\n    token_cell_indices = []\n    \n    for cur_cell_idx, code_token_ids in enumerate(code_token_ids):\n        token_count_for_cell = len(code_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        input_ids += code_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    \n    # Pad the features to match max_seq_len #\n    num_pad_tokens = max_seq_len-len(input_ids)\n    token_labels += [0]*num_pad_tokens\n    token_weights += [0]*num_pad_tokens\n    token_cell_indices += [-1]*num_pad_tokens\n    markdown_token_mask += [0]*num_pad_tokens\n    code_token_mask += [0]*num_pad_tokens\n    attention_mask = [1]*len(input_ids) + [0]*num_pad_tokens\n    input_ids += [0]*num_pad_tokens\n    \n    # Check for bugs\n    assert len(input_ids) == max_seq_len\n    assert len(token_labels) == max_seq_len\n    \n    # Build the feature dict for the input \n    notebook_features = {\n        'input_ids': input_ids, \n        'attention_mask': attention_mask,\n        'markdown_token_mask': markdown_token_mask,\n        'code_token_mask': code_token_mask,\n        'token_cell_indices': token_cell_indices,\n        'notebook_id': notebook_dict['notebook_id'],\n    }\n    if 'merged_cell_pct_ranks' in notebook_dict:\n        notebook_features['token_labels'] = token_labels\n        notebook_features['token_weights'] = token_weights\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_markdown_cell,\n    max_tokens_per_code_cell,\n    ):\n    '''Builds the huggingface dataset for training the model.'''\n    convert_to_features = partial(\n        convert_to_features_bigbird, \n        tokenizer=tokenizer,\n        max_seq_len=max_seq_len,\n        max_markdown_seq_len=max_markdown_seq_len,\n        max_tokens_per_markdown_cell=max_tokens_per_markdown_cell,\n        max_tokens_per_code_cell=max_tokens_per_code_cell,\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_length', default=1024, type=int, help='The max sequence length')\n    parser.add_argument('--max_markdown_seq_length', default=512, type=int, help='The max markdown sequence length')\n    parser.add_argument('--max_tokens_per_markdown_cell', default=128, type=int, help='The max tokens per markdown cell')\n    parser.add_argument('--max_tokens_per_code_cell', default=128, type=int, help='The max tokens per code cell')\n    parser.add_argument('--notebooks_df_path', default='notebooks_df.csv', type=str, help='Path to notebooks.csv')\n    parser.add_argument('--valid_fold', default=0, type=int, help='Validation Fold')\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    train_df = notebooks_df[notebooks_df.notebook_fold != args.valid_fold]\n    valid_df = notebooks_df[notebooks_df.notebook_fold == args.valid_fold]\n    \n    train_dataset = build_hf_dataset(\n        df=train_df,\n        tokenizer=tokenizer,\n        max_seq_len=args.max_seq_length,\n        max_markdown_seq_len=args.max_markdown_seq_length,\n        max_tokens_per_markdown_cell=args.max_tokens_per_markdown_cell,\n        max_tokens_per_code_cell=args.max_tokens_per_code_cell,\n    )\n    valid_dataset = build_hf_dataset(\n        df=valid_df,\n        tokenizer=tokenizer,\n        max_seq_len=args.max_seq_length,\n        max_markdown_seq_len=args.max_markdown_seq_length,\n        max_tokens_per_markdown_cell=args.max_tokens_per_markdown_cell,\n        max_tokens_per_code_cell=args.max_tokens_per_code_cell,\n    )\n    print('Train dataset size:', len(train_dataset))\n    print('Valid dataset size:', len(valid_dataset))\n\n    train_dataset.save_to_disk('train_dataset')\n    valid_dataset.save_to_disk('valid_dataset')\n    train_df.to_csv('train_df.csv', index=False)\n    valid_df.to_csv('valid_df.csv', index=False)\n    print('Done!')","metadata":{"execution":{"iopub.status.busy":"2022-07-14T21:53:35.917341Z","iopub.execute_input":"2022-07-14T21:53:35.917705Z","iopub.status.idle":"2022-07-14T21:53:35.985452Z","shell.execute_reply.started":"2022-07-14T21:53:35.917661Z","shell.execute_reply":"2022-07-14T21:53:35.984589Z"},"trusted":true},"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'fold{HP.valid_fold}_dataset')\n    train_hf_dataset = datasets.concatenate_datasets([\n        datasets.load_from_disk(processed_dataset_path/f'fold{fold}_dataset')\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_length, \n        max_markdown_seq_len=HP.max_markdown_seq_len,\n        max_tokens_per_markdown_cell=HP.max_tokens_per_markdown_cell, \n        max_tokens_per_code_cell=HP.max_tokens_per_code_cell,\n    )","metadata":{"execution":{"iopub.status.busy":"2022-07-14T21:53:35.986969Z","iopub.execute_input":"2022-07-14T21:53:35.987282Z","iopub.status.idle":"2022-07-14T21:54:51.428396Z","shell.execute_reply.started":"2022-07-14T21:53:35.987247Z","shell.execute_reply":"2022-07-14T21:54:51.427548Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 🌱 Flax Dataloader\n---\n\n#### <a href='#hyperparameters'> ⚙️ Hyperparameters </a>  | <a href='#model'> 🧠 Model </a> | <a href='#training-loop'> ⚡ Training Loop </a> \n","metadata":{}},{"cell_type":"code","source":"def collate_fn(hf_batch):\n    jax_batch = {\n        'input_ids': jnp.array(hf_batch['input_ids'], dtype=jnp.int32),\n        'attention_mask': jnp.array(hf_batch['attention_mask'], dtype=jnp.int32),\n        'markdown_token_mask': jnp.array(hf_batch['markdown_token_mask'], dtype=jnp.int32),\n        'code_token_mask': jnp.array(hf_batch['code_token_mask'], dtype=jnp.int32),\n        'token_labels': jnp.array(hf_batch['token_labels'], dtype=jnp.float32),\n        'token_weights': jnp.array(hf_batch['token_weights'], dtype=jnp.float32),\n    }\n    jax_batch = jax.tree_map(shard, jax_batch)\n    return jax_batch\n\ndef get_dataloader(hf_dataset, batch_size, shuffle_seed=None):\n    if shuffle_seed is not None:\n        hf_dataset = hf_dataset.shuffle(seed=shuffle_seed)\n    total_batches = len(hf_dataset)//batch_size\n    print(f'{total_batches} batches from {len(hf_dataset)} examples for dataloader')\n    for batch_idx in range(total_batches):\n        batch = hf_dataset[batch_idx * batch_size : (batch_idx + 1) * batch_size]\n        yield dict(batch)","metadata":{"execution":{"iopub.status.busy":"2022-07-14T21:54:51.430140Z","iopub.execute_input":"2022-07-14T21:54:51.430416Z","iopub.status.idle":"2022-07-14T21:54:51.520079Z","shell.execute_reply.started":"2022-07-14T21:54:51.430383Z","shell.execute_reply":"2022-07-14T21:54:51.519282Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 🐥 BigBird Model 🐥\n---\n### <a href='#hyperparameters'> ⚙️ Hyperparameters </a>  | <a href='#training-loop'> ⚡ Training Loop </a> \n\n<a name='model'>\n","metadata":{}},{"cell_type":"code","source":"model = transformers.FlaxBigBirdForTokenClassification.from_pretrained(\n    HP.backbone_name,\n    num_labels=1,\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    #num_random_blocks=HP.num_random_blocks,\n    #block_size=HP.block_size,\n    #dtype=jax.numpy.bfloat16,\n)","metadata":{"execution":{"iopub.status.busy":"2022-07-14T21:54:51.521203Z","iopub.execute_input":"2022-07-14T21:54:51.522211Z","iopub.status.idle":"2022-07-14T21:55:27.747330Z","shell.execute_reply.started":"2022-07-14T21:54:51.522164Z","shell.execute_reply":"2022-07-14T21:55:27.745752Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 🏭 Optimizer Factory\n---\n#### <a href='#hyperparameters'> ⚙️ Hyperparameters </a>  | <a href='#model'> 🧠 Model </a> | <a href='#training-loop'> ⚡ Training Loop </a> \n\n<a name='optimizer-factory'>","metadata":{}},{"cell_type":"code","source":"%%writefile optimizer_factory.py\n\nimport optax\nimport flax\n\ndef build_lr_scheduler(peak_lr, min_lr, warmup_ratio, total_train_steps):\n    warmup_steps = int(warmup_ratio * total_train_steps)\n    decay_steps = total_train_steps - warmup_steps\n    print(f'Warmup Steps: {warmup_steps} | Decay Steps: {decay_steps}')\n    \n    lr_scheduler = optax.warmup_cosine_decay_schedule(\n        init_value=min_lr,\n        peak_value=peak_lr,\n        warmup_steps=warmup_steps,\n        decay_steps=decay_steps,\n        end_value=min_lr,\n    )\n    return lr_scheduler\n\ndef build_tx(\n    lr_scheduler,\n    adam_beta_2,\n    adam_epsilon,\n    weight_decay,\n    max_grad_norm,\n    ema_decay,\n):\n    def weight_decay_mask(params):\n        params = flax.traverse_util.flatten_dict(params)\n        mask = {k: (v[-1] != 'bias' and v[-2:] != ('LayerNorm', 'scale')) for k, v in params.items()}\n        return flax.traverse_util.unflatten_dict(mask)\n\n    tx = optax.adamw(\n        learning_rate=lr_scheduler, \n        b1=0.9, \n        b2=adam_beta_2, \n        eps=adam_epsilon, \n        weight_decay=weight_decay, \n        mask=weight_decay_mask,\n    )\n\n    if max_grad_norm is not None:\n        tx = optax.chain(tx, optax.clip_by_global_norm(max_grad_norm))\n    if ema_decay is not None:\n        tx = optax.chain(tx, optax.ema(decay=ema_decay)) \n    return tx","metadata":{"execution":{"iopub.status.busy":"2022-07-14T21:55:27.749554Z","iopub.execute_input":"2022-07-14T21:55:27.749983Z","iopub.status.idle":"2022-07-14T21:55:27.844503Z","shell.execute_reply.started":"2022-07-14T21:55:27.749932Z","shell.execute_reply":"2022-07-14T21:55:27.842754Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import optimizer_factory\n\ntrain_batch_size = HP.per_device_train_batch_size * jax.device_count()\neval_batch_size = HP.per_device_eval_batch_size * jax.device_count()\ntotal_train_steps = HP.num_train_epochs * (len(train_hf_dataset) // train_batch_size)\n\nlr_scheduler = optimizer_factory.build_lr_scheduler(\n    peak_lr=HP.peak_lr,\n    min_lr=HP.min_lr,\n    warmup_ratio=HP.warmup_ratio,\n    total_train_steps=total_train_steps,\n)\n\ntx = optimizer_factory.build_tx(\n    lr_scheduler=lr_scheduler,\n    adam_beta_2=HP.beta_2,\n    adam_epsilon=HP.epsilon,\n    weight_decay=HP.weight_decay,\n    max_grad_norm=HP.max_grad_norm,\n    ema_decay=HP.ema_decay,\n)","metadata":{"execution":{"iopub.status.busy":"2022-07-14T21:55:27.848640Z","iopub.execute_input":"2022-07-14T21:55:27.849552Z","iopub.status.idle":"2022-07-14T21:55:27.950265Z","shell.execute_reply.started":"2022-07-14T21:55:27.849491Z","shell.execute_reply":"2022-07-14T21:55:27.949123Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## ⚡ Training Loop\n---\n### <a href='#hyperparameters'> ⚙️ Hyperparameters </a>  | <a href='#model'> 🐥 BigBird Model </a> \n\n<a name='training-loop'>","metadata":{}},{"cell_type":"code","source":"from typing import Callable\n\ndef compute_loss_per_cell_type(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\n    diff = jnp.abs(token_labels - token_preds)\n    notebook_token_weights_sum = jnp.sum(token_weights, axis=-1)\n\n    if HP.loss_fn_name == 'mse':\n        notebook_losses = jnp.sum(diff**2 * token_weights, axis=-1)\n        loss = jnp.sum(notebook_losses / notebook_token_weights_sum)\n    elif HP.loss_fn_name == 'rmse':\n        notebook_losses = jnp.sum(diff**2 * token_weights, axis=-1)\n        loss = jnp.sqrt(jnp.sum(notebook_losses / notebook_token_weights_sum))\n    elif HP.loss_fn_name == 'mae':\n        notebook_losses = jnp.sum(diff * token_weights, axis=-1)\n        loss = jnp.sum(notebook_losses / notebook_token_weights_sum)\n\n    return loss\n\ndef loss_fn(token_labels, token_preds, token_weights, markdown_token_mask, code_token_mask):\n    \n    markdown_cell_loss = compute_loss_per_cell_type(\n        token_labels=token_labels,\n        token_preds=token_preds,\n        token_weights=token_weights,\n        token_mask=markdown_token_mask,\n    )\n    code_cell_loss = compute_loss_per_cell_type(\n        token_labels=token_labels,\n        token_preds=token_preds,\n        token_weights=token_weights,\n        token_mask=code_token_mask,\n    )\n    return markdown_cell_loss, code_cell_loss\n\nclass TrainState(flax.training.train_state.TrainState):\n    loss_fn: Callable = flax.struct.field(pytree_node=False)\n    logits_fn: Callable = flax.struct.field(pytree_node=False)\n\ndef create_state(model, tx):\n    params = model.params\n    state = TrainState.create(\n        apply_fn=model.__call__,\n        params=params,\n        tx=tx,\n        loss_fn=loss_fn,\n        logits_fn=lambda x: x,\n    )\n    state = flax.jax_utils.replicate(state)\n    return state\n\n@partial(jax.pmap, axis_name='batch')\ndef train_step(state, drp_rng, **model_inputs):\n    def loss_fn(params):\n        outputs = state.apply_fn(\n            input_ids=model_inputs['input_ids'],\n            attention_mask=model_inputs['attention_mask'],\n            params=params,\n            dropout_rng=drp_rng,\n            train=True,\n        )\n        token_preds = outputs.logits\n        \n        token_preds = jnp.squeeze(token_preds)\n        token_labels = jnp.squeeze(model_inputs['token_labels'])\n        token_weights = jnp.squeeze(model_inputs['token_weights'])\n        markdown_token_mask = jnp.squeeze(model_inputs['markdown_token_mask'])\n        code_token_mask = jnp.squeeze(model_inputs['code_token_mask'])\n\n        markdown_cell_loss, code_cell_loss = state.loss_fn(\n            token_labels=token_labels,\n            token_preds=token_preds,\n            token_weights=token_weights,\n            markdown_token_mask=markdown_token_mask,\n            code_token_mask=code_token_mask,\n        )\n        \n        markdown_cell_loss_weight = HP.markdown_cell_loss_weight\n        code_cell_loss_weight = 1 - HP.markdown_cell_loss_weight\n        loss = markdown_cell_loss_weight * markdown_cell_loss + code_cell_loss_weight * code_cell_loss\n        \n        return loss, markdown_cell_loss\n\n    drp_rng, new_drop_rng = jax.random.split(drp_rng)\n    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n    (loss, markdown_cell_loss), grads = grad_fn(state.params)\n    \n    metrics = {\n        'loss': loss,\n        'markdown_cell_loss': markdown_cell_loss,\n        'learning_rate': lr_scheduler(state.step),\n    }\n    metrics = jax.lax.pmean(metrics, axis_name='batch')\n    grads = jax.lax.pmean(grads, 'batch')\n\n    state = state.apply_gradients(grads=grads)\n    return state, metrics, new_drop_rng\n\n\nprint('---------------- Available devices ----------------')\nprint(jax.devices())\nprint('---------------------------------------------------')\n\ntrain_steps_per_epoch = len(train_hf_dataset)//train_batch_size\ntotal_train_steps = train_steps_per_epoch * HP.num_train_epochs\nrng = jax.random.PRNGKey(HP.random_state)\ndrp_rng = jax.random.split(rng, jax.device_count())\nstate = create_state(model, tx)\n\nfor epoch in range(HP.num_train_epochs):\n    train_dataloader = get_dataloader(train_hf_dataset, train_batch_size, shuffle_seed=epoch)\n    running_metrics = defaultdict(int)\n    steps_progress_bar = tqdm(enumerate(train_dataloader), total=train_steps_per_epoch, desc=f'Epoch #{epoch+1}/{HP.num_train_epochs}')\n    \n    for step, hf_batch in steps_progress_bar:\n        batch = collate_fn(hf_batch)\n        state, step_metrics, drp_rng = train_step(\n            state=state,\n            drp_rng=drp_rng,\n            **batch\n        )\n        \n        # Update progress bar and running metrics\n        step_metrics = flax.jax_utils.unreplicate(step_metrics)\n        steps_progress_bar.set_postfix(**step_metrics)\n        for k, v in step_metrics.items():\n            running_metrics[k] += step_metrics[k]\n\n        # Log metrics every `logging_freq` steps\n        if (step + 1) % HP.logging_freq == 0:\n            global_step = flax.jax_utils.unreplicate(state.step)\n            print('-'*50)\n            print(f\"Step {global_step-HP.logging_freq}-{global_step} out of {total_train_steps}\")\n            for k, v in running_metrics.items():\n                print(colored(k, 'blue'), ':', colored(v / HP.logging_freq, 'red'))\n            running_metrics = defaultdict(int)\n            print()\n\n    print(colored(f'Epoch #{epoch+1} completed.'))\n    print(colored('-'*100, 'red'))\n    print('\\n\\n')","metadata":{"execution":{"iopub.status.busy":"2022-07-15T00:08:54.982377Z","iopub.execute_input":"2022-07-15T00:08:54.982949Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 🎯 Inference\n---\n\n<a name='inference'>","metadata":{}},{"cell_type":"code","source":"# CELL_SEP = '[CELL_SEP]'\n\n# @partial(jax.pmap, axis_name='batch')\n# def predict_step(state, batch):\n#     input_ids = jnp.squeeze(batch['input_ids'])\n#     attention_mask = jnp.squeeze(batch['attention_mask'])\n    \n#     outputs = state.apply_fn(\n#         input_ids=input_ids,\n#         attention_mask=attention_mask,\n#         params=state.params,\n#         train=False,\n#     )\n    \n#     token_preds = jnp.squeeze(outputs.logits)\n#     token_labels = jnp.squeeze(batch['token_labels'])\n#     token_weights = jnp.squeeze(batch['token_weights'])\n#     markdown_token_mask = jnp.squeeze(batch['markdown_token_mask'])\n#     code_token_mask = jnp.squeeze(batch['code_token_mask'])\n#     markdown_cell_loss, code_cell_loss = state.loss_fn(\n#         token_labels=token_labels,\n#         token_preds=token_preds,\n#         token_weights=token_weights,\n#         markdown_token_mask=markdown_token_mask,\n#         code_token_mask=code_token_mask,\n#     )\n    \n#     metrics = {\n#         'markdown_cell_loss': markdown_cell_loss, \n#         'code_cell_loss': code_cell_loss\n#     }\n#     metrics = jax.lax.pmean(metrics, axis_name=\"batch\")\n#     return metrics\n\n# valid_dataloader = get_dataloader(valid_hf_dataset, eval_batch_size, shuffle_seed=epoch)\n# agg_metrics = defaultdict(list)\n# for hf_batch in tqdm(valid_dataloader, total=len(valid_hf_dataset)//eval_batch_size):\n#     batch = collate_fn(hf_batch)\n#     step_metrics = predict_step(state, batch)\n#     for k, v in step_metrics.items():\n#         agg_metrics[k].append(v)\n# for k, v in agg_metrics.items():\n#     v = np.mean(np.array(v))\n#     print(colored(k, 'red'), ':', colored(v, 'blue'))","metadata":{"execution":{"iopub.status.busy":"2022-07-15T00:04:43.646214Z","iopub.status.idle":"2022-07-15T00:04:43.647219Z","shell.execute_reply.started":"2022-07-15T00:04:43.646905Z","shell.execute_reply":"2022-07-15T00:04:43.646937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# CELL_SEP = '[CELL_SEP]'\n\n# @partial(jax.pmap, axis_name='batch')\n# def predict_step(state, batch):\n2\n#     return state.logits_fn(token_preds)\n    \n#     #return jnp.asarray(token_preds*10, dtype=jnp.int8)\n\n# def compute_kendall_tau(valid_df, cell_id_to_pred_rank):\n#     all_notebook_cell_pct_ranks = valid_df.merged_cell_pct_ranks.values\n#     all_notebook_cell_ids = valid_df.merged_cell_ids.values\n#     all_notebook_kendall_taus = []\n#     for 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, -1) for cell_id in cell_ids]\n\n#         notebook_kendall_tau = scipy.stats.kendalltau(true_cell_ranks, pred_cell_ranks, method='asymptotic')[0]\n#         all_notebook_kendall_taus.append(notebook_kendall_tau)\n\n#     valid_df['kendall_tau'] = all_notebook_kendall_taus\n#     print('Average Kendall Tau:', colored(np.mean(all_notebook_kendall_taus), 'red'))\n\n#     for cell_cutoff in [4, 16, 64]:\n#         tau = np.mean(valid_df[valid_df.markdown_cell_count >= cell_cutoff].kendall_tau)\n#         print(f\"Kendall Tau for notebooks with {cell_cutoff}+ markdown cells:\", colored(tau, 'red'))\n\n# cell_id_to_token_preds = defaultdict(list)\n# valid_dataloader = get_dataloader(valid_hf_dataset, eval_batch_size, shuffle_seed=epoch)\n# for hf_batch in tqdm(valid_dataloader, total=len(valid_hf_dataset)//eval_batch_size):\n#     batch = collate_fn(hf_batch)\n#     token_preds = predict_step(state, batch)\n#     token_preds = np.array([pred for pred in itertools.chain(*token_preds)])\n    \n#     token_cell_indices = hf_batch['token_cell_indices']\n#     notebook_ids = hf_batch['notebook_id']\n#     for example_idx, notebook_id in enumerate(notebook_ids):\n#         cell_ids = valid_df[valid_df.notebook_id==notebook_id].iloc[0].merged_cell_ids.split(CELL_SEP)\n#         for cell_idx, token_pred in zip(token_cell_indices[example_idx], token_preds[example_idx]):\n#             cell_id = cell_ids[cell_idx]\n#             cell_id_to_token_preds[cell_id].append(token_pred)\n# cell_id_to_pred_rank = {cell_id: np.mean(preds) for cell_id, preds in cell_id_to_token_preds.items()}\n# compute_kendall_tau(valid_df, cell_id_to_pred_rank)\n# print('\\n\\n')","metadata":{"execution":{"iopub.status.busy":"2022-07-15T00:04:43.649493Z","iopub.status.idle":"2022-07-15T00:04:43.650265Z","shell.execute_reply.started":"2022-07-15T00:04:43.649965Z","shell.execute_reply":"2022-07-15T00:04:43.650007Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# avg_tau = valid_df.kendall_tau.mean()\n# tau_16 = np.mean(valid_df[valid_df.markdown_cell_count >= 16].kendall_tau)\n# tau_64 = np.mean(valid_df[valid_df.markdown_cell_count >= 64].kendall_tau)\n\n# WEIGHTS_SAVE_FORMAT = \"{backbone_code}-tau{avg_tau}-tau16_{tau_16}-tau64_{tau_64}.msgpack\"\n# weights_file = WEIGHTS_SAVE_FORMAT.format(\n#     backbone_code=backbone_code,\n#     avg_tau=int(avg_tau*1e6),\n#     tau_16=int(tau_16*1e4),\n#     tau_64=int(tau_64*1e4),\n# )\n# print(f'Saving weights at {weights_file}')\n# if jax.process_index() == 0:\n#     params = jax.device_get(flax.jax_utils.unreplicate(state.params))\n    \n#     with open(f'/kaggle/working/{weights_file}', 'wb') as f:\n#         model_bytes = flax.serialization.to_bytes(params)\n#         f.write(model_bytes)","metadata":{"execution":{"iopub.status.busy":"2022-07-15T00:04:43.651622Z","iopub.status.idle":"2022-07-15T00:04:43.653076Z","shell.execute_reply.started":"2022-07-15T00:04:43.652550Z","shell.execute_reply":"2022-07-15T00:04:43.652611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"weights_file = 'bigbird-seq1536_md1024-loss146.msgpack'\nprint(weights_file)\nif jax.process_index() == 0:\n    params = jax.device_get(flax.jax_utils.unreplicate(state.params))\n    with open(f'/kaggle/working/{weights_file}', 'wb') as f:\n        model_bytes = flax.serialization.to_bytes(params)\n        f.write(model_bytes)","metadata":{"execution":{"iopub.status.busy":"2022-07-15T00:05:36.344925Z","iopub.execute_input":"2022-07-15T00:05:36.345487Z","iopub.status.idle":"2022-07-15T00:05:39.438123Z","shell.execute_reply.started":"2022-07-15T00:05:36.345437Z","shell.execute_reply":"2022-07-15T00:05:39.436753Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"wandb.init(project='bigbird_dev')\nwandb.save(weights_file)\nsleep(180)","metadata":{"execution":{"iopub.status.busy":"2022-07-15T00:05:43.030202Z","iopub.execute_input":"2022-07-15T00:05:43.030631Z","iopub.status.idle":"2022-07-15T00:08:54.979517Z","shell.execute_reply.started":"2022-07-15T00:05:43.030577Z","shell.execute_reply":"2022-07-15T00:08:54.978334Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}