{"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 TFRankEncoder Train (CV: 0.85+)__\n\n---\n### <a href='#hyperparameters'> ⚙️ Hyperparameters </a> | <a href='#data-factory'> ⚒ Data Factory </a>  | <a href='#training'> ⚡ Training </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":"# Sync Notebook with VS Code #\n!pip install -q transformers==4.10.0 datasets\nimport sys; sys.path.append('ai4code')\n\n%run /kaggle/working/ai4code/ai4c/jupyter_setup.py\n%run /kaggle/working/ai4code/ai4c/tensorflow_setup.py","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":42.720592,"end_time":"2022-05-21T12:18:38.820861","exception":false,"start_time":"2022-05-21T12:17:56.100269","status":"completed"},"tags":[],"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'> ⚡ Training </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\n\n## Huggingface Backbone ##\nbackbone_name: 'facebook/muppet-roberta-large'\nbackbone_weights: null\n\nattention_probs_dropout_prob: 0.10\nhidden_dropout_prob: 0.10\nmax_seq_len: 512\n\n\n## Tokenization ##\nmax_markdown_seq_len: 256\nmax_tokens_per_markdown_cell: 128\nmax_tokens_per_code_cell: 128\n\n\n## Model Training ##\nnum_train_epochs: 10\ntrain_batch_size: 256\neval_batch_size: 1024\n\n\n## Loss Function ##\nloss_fn: 'mse'\nmarkdown_cell_weight: 0.50\n\n\n## Cosine Decay LR Scheduler ##\nwarmup_ratio: 0.10\npeak_lr: 3e-5\nmin_lr: 1e-8\n\n\n## AdamW Optimizer ##\nbeta_1: 0.9\nbeta_2: 0.999\nepsilon: 1e-8\n\nweight_decay: 1e-4\nmax_grad_norm: 1.00\naverage_decay: 0.999\n\n\n## Load From Cache: Tokenized Dataset ##\nprocessed_dataset_folder: null # 'ai4code-tokenization-bigbird'\ndebug_notebooks: 10000\n\n\n## Data Factory ##\nvalid_fold: 0","metadata":{"papermill":{"duration":0.112836,"end_time":"2022-05-21T12:18:38.990709","exception":false,"start_time":"2022-05-21T12:18:38.877873","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"STRATEGY = tf_accelerator(bfloat16=True, jit_compile=True)\nwith STRATEGY.scope():\n    backbone = TFAutoModel.from_pretrained(\n        HP.backbone_name, \n        attention_probs_dropout_prob=HP.attention_probs_dropout_prob,\n        hidden_dropout_prob=HP.hidden_dropout_prob,\n        from_pt=True\n    )\n    load_tf_model_weights(backbone, HP.backbone_weights)\n\ntokenizer = AutoTokenizer.from_pretrained(HP.backbone_name)","metadata":{"papermill":{"duration":119.07852,"end_time":"2022-05-21T12:20:38.333669","exception":false,"start_time":"2022-05-21T12:18:39.255149","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## ⚒️ Data Factory ⚒️\n\n---\n#### <a href='#prepare-huggingface-datasets'> 🤗 Huggingface Datasets </a> | <a href='#prepare-tensorflow-datasets'> Tensorflow Datasets </a> \n\n\n<a name='data-factory'>","metadata":{}},{"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]\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":{"trusted":true},"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":{}},{"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":{"papermill":{"duration":1.10202,"end_time":"2022-05-21T12:23:23.649546","exception":false,"start_time":"2022-05-21T12:23:22.547526","status":"completed"},"tags":[],"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_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Prepare Tensorflow 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-tensorflow-datasets'>","metadata":{}},{"cell_type":"code","source":"%%writefile prepare_tf_data.py\n\ndef convert_hf_dataset_to_tfds(hf_dataset):\n    hf_dataset.set_format(type='numpy')\n\n    model_inputs = {\n        'input_ids': hf_dataset['input_ids'].astype(np.int32),\n        'attention_mask': hf_dataset['attention_mask'].astype(np.int32),\n        'markdown_token_mask': hf_dataset['markdown_token_mask'].astype(np.int32),\n        'code_token_mask': hf_dataset['code_token_mask'].astype(np.int32),\n    }\n    input_ds = tf.data.Dataset.from_tensor_slices(model_inputs)\n\n    model_outputs = {\n        'token_labels': hf_dataset['token_labels'].astype(np.float32),\n    }\n    output_ds = tf.data.Dataset.from_tensor_slices(model_outputs)\n\n    ds = tf.data.Dataset.zip((input_ds, output_ds))\n    return ds\n\ndef hf_dataset_to_tfds(hf_dataset, dataset_type, batch_size): \n    ds = convert_hf_dataset_to_tfds(hf_dataset)\n    if dataset_type == 'train':\n        ds = ds.shuffle(len(hf_dataset), reshuffle_each_iteration=True).repeat()\n    elif dataset_type == 'valid': \n        ds = ds.cache()\n    ds = ds.batch(batch_size)\n    steps = len(hf_dataset)//batch_size + 1\n    return ds.prefetch(tf.data.AUTOTUNE), steps\n\ndef convert_hf_dataset_to_test_ds(hf_dataset, batch_size):\n    input_ids_ds = tf.data.Dataset.from_tensor_slices(hf_dataset['input_id'].astype(np.int32))\n    mask_ds = tf.data.Dataset.from_tensor_slices(hf_dataset['attention_mask'].astype(np.int32))\n    input_ds = tf.data.Dataset.zip((input_ids_ds, mask_ds))\n    ds = tf.data.Dataset.zip((input_ds, input_ds))\n    return ds.batch(batch_size).prefetch(tf.data.AUTOTUNE)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import prepare_tf_data\n\ntrain_ds, train_steps_per_epoch = prepare_tf_data.hf_dataset_to_tfds(\n    hf_dataset=train_hf_dataset,\n    dataset_type='train',\n    batch_size=HP.train_batch_size\n)\nvalid_ds, valid_steps_per_epoch = prepare_tf_data.hf_dataset_to_tfds(\n    hf_dataset=valid_hf_dataset,\n    dataset_type='valid',\n    batch_size=HP.eval_batch_size,\n)\n\neval_ds = prepare_tf_data.convert_hf_dataset_to_test_ds(valid_hf_dataset, HP.eval_batch_size)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 🧠 Model Factory\n---\n#### <a href='#training'> ⚡ Training </a>\n\n<a name='model-factory'>","metadata":{"papermill":{"duration":0.027668,"end_time":"2022-05-21T12:23:23.705434","exception":false,"start_time":"2022-05-21T12:23:23.677766","status":"completed"},"tags":[]}},{"cell_type":"code","source":"%%writefile tf_rankencoder_model.py\nimport tensorflow as tf\n\nclass AI4CodeTFRankEncoder(tf.keras.Model):\n    def __init__(\n        self, \n        model_inputs, \n        model_outputs,\n        markdown_cell_loss_weight=0.50,\n        loss_fn_name='mse',\n    ):\n        super().__init__(inputs=model_inputs, outputs=model_outputs)\n        self.metrics_tracker = {\n            'total_loss': tf.keras.metrics.Mean(name='loss'),\n            'markdown_cell_loss': tf.keras.metrics.Mean(name='markdown_cell_loss'),\n            'code_cell_loss': tf.keras.metrics.Mean(name='code_cell_loss'),\n            'gradient_norm': tf.keras.metrics.Mean(name='gradient_norm'),\n        }\n        self.markdown_cell_loss_weight = markdown_cell_loss_weight\n        self.code_cell_loss_weight = 1.0 - markdown_cell_loss_weight\n\n        self.loss_fn = {\n            'mse': self.mse_loss_fn,\n            'mae': self.mae_loss_fn,\n        }[loss_fn_name]\n    \n    def mse_loss_fn(self, y_true, y_pred, token_mask):\n        diff = (y_true-y_pred)*token_mask\n        return tf.math.reduce_sum(diff**2) / tf.math.reduce_sum(token_mask)\n    \n    def mae_loss_fn(self, y_true, y_pred, token_mask):\n        diff = (y_true-y_pred)*token_mask\n        return tf.math.reduce_sum(tf.math.abs(diff)) / tf.math.reduce_sum(token_mask)\n    \n    @tf.function\n    def train_step(self, data):\n        x, y = data\n        with tf.GradientTape() as tape: \n            model_inputs = (x['input_ids'], x['attention_mask'])\n            token_labels_pred = self(model_inputs)\n            \n            token_labels_pred = tf.cast(token_labels_pred, tf.float32)\n            token_labels_true = tf.cast(y['token_labels'], tf.float32)\n            markdown_cell_mask = tf.cast(x['markdown_cell_mask'], tf.float32)\n            code_cell_mask = tf.cast(x['code_cell_mask'], tf.float32)\n\n            markdown_cell_loss = self.loss_fn(token_labels_true, token_labels_pred, markdown_cell_mask)\n            code_cell_loss = self.loss_fn(token_labels_true, token_labels_pred, code_cell_mask)\n            \n            loss = self.markdown_cell_weight*markdown_cell_loss + self.code_cell_weight*code_cell_loss\n        \n        trainable_vars = self.trainable_variables\n        gradients = tape.gradient(loss, trainable_vars)\n        gradient_norm = tf.linalg.global_norm(gradients)\n        self.optimizer.apply_gradients(zip(gradients, trainable_vars))\n        \n        self.metrics_tracker['total_loss'].update_state(loss)\n        self.metrics_tracker['code_cell_loss'].update_state(code_cell_loss)\n        self.metrics_tracker['markdown_cell_loss'].update_state(markdown_cell_loss)\n        self.metrics_tracker['gradient_norm'].update_state(gradient_norm)\n        return {m.name: m.result() for m in self.metrics}\n    \n    def test_step(self, data):\n        x, y = data\n        model_inputs = (x['input_ids'], x['attention_mask'])\n        token_labels_pred = self(model_inputs)\n        \n        token_labels_pred = tf.cast(token_labels_pred, tf.float32)\n        token_labels_true = tf.cast(y['token_labels'], tf.float32)\n        markdown_cell_mask = tf.cast(x['markdown_cell_mask'], tf.float32)\n        code_cell_mask = tf.cast(x['code_cell_mask'], tf.float32)\n\n        markdown_cell_loss = self.loss_fn(token_labels_true, token_labels_pred, markdown_cell_mask)\n        code_cell_loss = self.loss_fn(token_labels_true, token_labels_pred, code_cell_mask)\n        \n        loss = self.markdown_cell_weight*markdown_cell_loss + self.code_cell_weight*code_cell_loss\n        \n        self.metrics_tracker['total_loss'].update_state(loss)\n        self.metrics_tracker['code_cell_loss'].update_state(code_cell_loss)\n        self.metrics_tracker['markdown_cell_loss'].update_state(markdown_cell_loss)\n        return {m.name: m.result() for m in self.metrics}\n    \n    @property\n    def metrics(self):\n        return list(self.metrics_tracker.values())\n\ndef build_tfrankencoder_model(\n    backbone, \n    max_seq_len,\n    markdown_cell_loss_weight=0.50,\n    loss_fn_name='mse',\n    ):\n\n    input_ids = tf.keras.Input(shape=(max_seq_len,), dtype=tf.int32, name='input_ids')\n    attention_mask = tf.keras.Input(shape=(max_seq_len,), dtype=tf.float32, name='attention_mask')\n    model_inputs = [input_ids, attention_mask]\n\n    token_ranker_layer = tf.keras.Sequential([\n        tf.keras.layers.Dense(1), \n        tf.keras.layers.Reshape((max_seq_len,))\n    ], name='token_labels')\n\n    backbone_outputs = backbone(input_ids=input_ids, attention_mask=attention_mask)\n    x = backbone_outputs.last_hidden_state\n    return AI4CodeTFRankEncoder(\n        model_inputs=model_inputs,\n        model_outputs=token_ranker_layer(x),\n        markdown_cell_loss_weight=markdown_cell_loss_weight,\n        loss_fn_name=loss_fn_name,\n    )","metadata":{"papermill":{"duration":10.462645,"end_time":"2022-05-21T12:23:34.345642","exception":false,"start_time":"2022-05-21T12:23:23.882997","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def adamw_optimizer_factory(lr_scheduler): \n    optimizer = tfa.optimizers.AdamW(\n        beta_1=HP.beta_1, \n        beta_2=HP.beta_2, \n        epsilon=HP.epsilon, \n        weight_decay=HP.max_weight_decay, \n        clipnorm=HP.max_grad_norm,\n        learning_rate=lr_scheduler,\n    )\n    if HP.average_decay > 0: \n        print(f\"Using EMA with decay {colored(HP.average_decay, 'blue')}\")\n        optimizer = tfa.optimizers.MovingAverage(\n            optimizer, \n            average_decay=HP.average_decay, \n            dynamic_decay=True, \n        )\n    return optimizer","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tf_rankencoder_model\n\ntotal_train_steps = train_steps_per_epoch * HP.num_train_epochs\nwarmup_steps = int(total_train_steps * HP.warmup_ratio)\ndecay_steps = total_train_steps - warmup_steps\n\ndecay_schedule_fn = tf.keras.optimizers.schedules.CosineDecay(\n    initial_learning_rate=HP.peak_lr, \n    decay_steps=decay_steps, \n    alpha=HP.min_lr,\n)\nlr_scheduler = transformers.WarmUp(\n    initial_learning_rate=HP.min_lr,\n    decay_schedule_fn=decay_schedule_fn,\n    warmup_steps=warmup_steps,\n)\n\nwith STRATEGY.scope():\n    model = tf_rankencoder_model.build_tfrankencoder_model(\n        backbone=backbone,\n        max_seq_len=HP.max_seq_len,\n        markdown_cell_loss_weight=HP.markdown_cell_loss_weight,\n        loss_fn_name=HP.loss_fn_name,               \n    )","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from collections import defaultdict\nfrom termcolor import colored\nfrom tqdm.auto import tqdm\nimport numpy as np\nimport scipy.stats\nimport os\n\nimport tensorflow as tf\n\nCELL_SEP = '[CELL_SEP]'\n\nclass KendallTauCallback(tf.keras.callbacks.Callback):\n    \"\"\"\n    Evaluates Kendall Tau for the notebooks at the end of each epoch\n    \"\"\"\n\n    def __init__(self, eval_df, eval_dataset, eval_ds, strategy):\n        self.eval_df = eval_df\n        self.eval_dataset = eval_dataset\n        self.eval_ds = eval_ds\n        self.strategy = strategy\n\n        self.saved_weights = []\n        os.makedirs('/kaggle/tmp', exist_ok=True)\n    \n    def evaluate_model(self, model):\n        with self.strategy.scope():\n            all_notebook_token_preds = np.array(model.predict(self.eval_ds, verbose=True)).astype(np.float32)\n        all_notebook_token_cell_indices = self.eval_dataset['token_cell_indices'].values\n\n        notebook_kendall_taus = []\n        for notebook_idx, (token_preds, token_cell_indices) in tqdm(\n            enumerate(zip(all_notebook_token_preds, all_notebook_token_cell_indices)), \n            total=len(self.eval_dataset)\n        ):\n            cell_idx_to_pred = defaultdict(list)\n            for cell_idx, token_pred in zip(token_cell_indices, token_preds):\n                cell_idx_to_pred[cell_idx].append(token_pred)\n            cell_idx_to_pred = {cell_idx: np.mean(preds) for cell_idx, preds in cell_idx_to_pred.items()}\n\n            notebook_cell_pct_ranks = [float(rank) for rank in self.eval_df.iloc[notebook_idx].merged_cell_pct_ranks.split(CELL_SEP)]\n            notebook_cell_preds = [cell_idx_to_pred[cell_idx] for cell_idx in range(len(notebook_cell_pct_ranks))]          \n            notebook_tau = scipy.stats.kendalltau(notebook_cell_pct_ranks, notebook_cell_preds, method='asymptotic')[0]\n            notebook_kendall_taus.append(notebook_tau)\n    \n        self.eval_df['kendall_tau'] = notebook_kendall_taus\n        for cell_cutoff in [4, 16, 64]:\n            cutoff_df = self.eval_df[self.eval_df.markdown_cell_count > cell_cutoff]\n            tau = np.mean(cutoff_df.kendall_tau.values)\n            print(f\"Average Kendall Tau for notebooks with {cell_cutoff}+ markdown cells: {colored(tau, 'red')}\")\n        \n        avg_kendall_tau = np.mean(notebook_kendall_taus)\n        print('Average Kendall Tau:', colored(avg_kendall_tau, 'red'))\n        return avg_kendall_tau\n\n    def on_epoch_end(self, epoch, logs=None):\n        epoch_kendall_tau = self.evaluate_model(self.model)\n        weights_file = f'/kaggle/tmp/weights_epoch_{epoch}_tau{int(epoch_kendall_tau*1e6)}.h5'\n        print('Saving weights at', colored(weights_file, 'blue'))\n        self.model.save_weights(weights_file)\n        self.saved_weights.append(weights_file)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"kendall_tau_callback = KendallTauCallback(\n    eval_df=valid_df,\n    eval_dataset=valid_dataset,\n    eval_ds=eval_ds,\n    strategy=STRATEGY,\n)","metadata":{},"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.028,"end_time":"2022-05-21T12:23:34.40338","exception":false,"start_time":"2022-05-21T12:23:34.37538","status":"completed"},"tags":[]}},{"cell_type":"code","source":"HP.multi_steps_per_execution = None\nwith STRATEGY.scope():\n    model = tf_rankencoder_model.build_tfrankencoder_model(\n        backbone=backbone,\n        max_seq_len=HP.max_seq_len,\n        markdown_cell_loss_weight=HP.markdown_cell_loss_weight,\n        loss_fn_name=HP.loss_fn_name,               \n    )\n    optimizer = adamw_optimizer_factory(lr_scheduler)\n    model.compile(\n        optimizer=optimizer, \n        steps_per_execution=HP.multi_steps_per_execution,\n    )\n\nhistory = model.fit(\n    train_ds, steps_per_epoch=train_steps_per_epoch, epochs=HP.num_train_epochs,\n    validation_data=valid_ds, validation_steps=valid_steps_per_epoch,\n    callbacks=[kendall_tau_callback]\n)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"HP.multi_steps_per_execution = False\nwith STRATEGY.scope():\n    model = build_model(backbone)\n    model = get_freeze_compiled_model()\n\nhistory = model.fit(train_ds, steps_per_epoch=train_steps)\ntau_callback.evaluate_model(model)\nmodel.save_weights('inital_model_weights.h5')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"HP.multi_steps_per_execution = True\nwith STRATEGY.scope():\n    backbone.trainable = True\n    model = build_model(backbone)\n    model = get_compiled_model()\nmodel.load_weights('inital_model_weights.h5')\n\n# 653s + 421s/epoch\nhistory = model.fit(\n    train_ds, steps_per_epoch=train_steps*4, epochs=HP.num_epochs,\n    validation_data=valid_ds, validation_steps=valid_steps,\n    callbacks=[tau_callback],\n)\ntf.keras.backend.clear_session()\ngc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_weight_files = sorted(tau_callback.saved_weights_files)[-HP.num_swa_models:]\nprint('Taking average of model weights:', model_weight_files)\nwith STRATEGY.scope():\n    model = get_model_average(model, model_weight_files)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# wandb.init(project='ai4code_tfrankencoder_fold0')\n\n# Save the model weights #\ntau = tau_callback.evaluate_model(model)\nweights_file = HP.wandb_weights_save_format.format(fold=HP.valid_fold, backbone_code=backbone_code, tau=int(tau*1000000))\nprint(blue(weights_file))\nmodel.save_weights(weights_file)\nwandb.save(weights_file)\n\n# 👽 Save the train and validation file #\ntrain.to_csv('train.csv', index=False)\nvalid.to_csv('valid.csv', index=False)\n\nsleep(180)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 👽 Paranoid Validation\n---\n\n<a name='paranoid-validation'>","metadata":{}},{"cell_type":"code","source":"cell_df = pd.read_csv('df.csv').set_index('notebook_id')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"i = 910\nrow = valid.iloc[i]\nnotebook_id = row.notebook_id\n\nprint('notebook id:', notebook_id)\nprint('notebook tau:', row.tau)\n\ncell_pct_preds = [float(pred) for pred in row.cell_pct_preds.split(CELL_SEP)]\ndisplay(cell_df.loc[notebook_id].sort_values(by='cell_pct_rank'))\n\ninput_ids, token_labels = valid_dataset[i]['input_ids'], valid_dataset[i]['token_labels']\nprint(tokenizer.decode(input_ids))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cell_idx, total_loss = 0, 0.0\nfor token_idx, (token_id, token_label) in enumerate(zip(input_ids, token_labels)):\n    if token_label == -100:\n        continue\n    print(red('-')*100)\n    print('Label:'+yellow(token_label))\n    print('Prediction:'+blue(cell_pct_preds[cell_idx]))\n    print()\n    for next_token_idx, (token_id, next_token_label) in enumerate(zip(input_ids, token_labels)):\n        if next_token_label == -100 or next_token_idx <= token_idx:\n            continue\n        break\n    print(tokenizer.decode(input_ids[token_idx:next_token_idx]))\n    total_loss += (cell_pct_preds[cell_idx]-token_label)**2\n    cell_idx += 1\n\nprint('Tau for the notebook:', row.tau)\nprint('MSE loss for notebook:', red(total_loss/cell_idx))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# # Basic Validation \n# valid, valid_dataset = train, train_dataset\n# eval_ds = convert_dataset_to_test_ds(valid_dataset)\n# tau_callback.evaluate_model(model)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}