{"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":"<img src=\"https://storage.googleapis.com/kaggle-competitions/kaggle/35887/logos/header.png?t=2022-05-09-22-33-02\">\n\n<h1><center>[3/3] AI4Code TensorFlow TPU with CodeBert - Inference</center></h1>\n\nThis is the final part of my **AI4Code TensorFlow TPU with CodeBert** series:\n\n* [1/3] [Data Preparation][1] (~5 hours)\n* [2/3] [TPU Training][2] (~4 hours)\n* **[3/3] GPU Inference ← (you're here)**\n\nThis is basically a translation of **[Khoi Nguyen's][3]** works [[1][4], [2][5]] from PyTorch to TensorFlow with minor changes and updates for TPU support. The **[original][4]** PyTorch work takes up to 40 hours per epoch on Kaggle GPU, whereas **[my version][2]** takes only 50 minutes per epoch on Kaggle TPU, so it's lightning fast ⚡.\n\n### About Solution\n\n- Input data: markdown + code context (512 tokens) + features\n    - Markdown (up to 64 tokens)\n    - Code context (all code cells or up to 20 code cells each up to 23 tokens)\n    - Features: markdown cells to total cells ratio (appended to backbone outputs)\n- Model and hyperparameters\n    - CodeBert Base model\n    - L1 loss (MAE)\n    - AdamW optimizer\n    - Learning rate schedule with warmup and linear decay\n    - Total 5 epochs\n\n### Input Data\n\n- **[AI4Code-CodeBert-Weights][6]**: output from **[TPU Training][2]** step\n- **[codebert-base][7]**: `microsoft/codebert-base` model saved locally (workaround for turned off internet requirement)\n\n[1]: https://www.kaggle.com/nickuzmenkov/ai4code-tf-tpu-codebert-data-preparation\n[2]: https://www.kaggle.com/nickuzmenkov/ai4code-tf-tpu-codebert-training\n[3]: https://www.kaggle.com/suicaokhoailang\n[4]: https://github.com/suicao/ai4code-baseline/tree/main/code\n[5]: https://www.kaggle.com/code/suicaokhoailang/stronger-baseline-with-code-cells\n[6]: https://www.kaggle.com/datasets/nickuzmenkov/ai4code-codebert-weights\n[7]: https://www.kaggle.com/datasets/leolu1998/codebert-base\n\n# Setup","metadata":{}},{"cell_type":"code","source":"import glob\nimport os\nfrom typing import List\n\nimport numpy as np\nimport pandas as pd\nimport tensorflow as tf\nimport transformers\nfrom tqdm.notebook import tqdm","metadata":{"execution":{"iopub.status.busy":"2022-06-22T14:57:47.31088Z","iopub.execute_input":"2022-06-22T14:57:47.311338Z","iopub.status.idle":"2022-06-22T14:57:51.834559Z","shell.execute_reply.started":"2022-06-22T14:57:47.311245Z","shell.execute_reply":"2022-06-22T14:57:51.83353Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BATCH_SIZE = 32\nSLICES = 8\nMD_MAX_LEN = 64\nTOTAL_MAX_LEN = 512\nSTRATEGY = tf.distribute.get_strategy()\nBASE_MODEL = \"../input/codebert-base/codebert-base\"\nTOKENIZER = transformers.AutoTokenizer.from_pretrained(BASE_MODEL)\nINPUT_PATH = \"../input/AI4Code\"","metadata":{"execution":{"iopub.status.busy":"2022-06-22T14:57:51.836927Z","iopub.execute_input":"2022-06-22T14:57:51.837619Z","iopub.status.idle":"2022-06-22T14:57:52.267774Z","shell.execute_reply.started":"2022-06-22T14:57:51.837581Z","shell.execute_reply":"2022-06-22T14:57:52.266788Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def read_notebook(path: str) -> pd.DataFrame:\n    return (\n        pd.read_json(path, dtype={\"cell_type\": \"category\", \"source\": \"str\"})\n        .assign(id=os.path.basename(path).split(\".\")[0])\n        .rename_axis(\"cell_id\")\n    )\n\n\ndef clean_code(cell: str) -> str:\n    return str(cell).replace(\"\\\\n\", \"\\n\")\n\n\ndef sample_cells(cells: List[str], n: int) -> List[str]:\n    cells = [clean_code(cell) for cell in cells]\n    if n >= len(cells):\n        return cells\n    else:\n        results = []\n        step = len(cells) / n\n        idx = 0\n        while int(np.round(idx)) < len(cells):\n            results.append(cells[int(np.round(idx))])\n            idx += step\n        if cells[-1] not in results:\n            results[-1] = cells[-1]\n        return results\n\n\ndef get_features(df: pd.DataFrame) -> dict:\n    features = {}\n    for i, sub_df in tqdm(df.groupby(\"id\"), desc=\"Features\"):\n        features[i] = {}\n        total_md = sub_df[sub_df.cell_type == \"markdown\"].shape[0]\n        code_sub_df = sub_df[sub_df.cell_type == \"code\"]\n        total_code = code_sub_df.shape[0]\n        codes = sample_cells(code_sub_df.source.values, 20)\n        features[i][\"total_code\"] = total_code\n        features[i][\"total_md\"] = total_md\n        features[i][\"codes\"] = codes\n    return features\n\n\ndef tokenize(df: pd.DataFrame, fts: dict) -> dict:\n    input_ids = np.zeros((len(df), TOTAL_MAX_LEN), dtype=np.int32)\n    attention_mask = np.zeros((len(df), TOTAL_MAX_LEN), dtype=np.int32)\n    features = np.zeros((len(df),), dtype=np.float32)\n\n    for i, row in tqdm(\n        df.reset_index(drop=True).iterrows(), desc=\"Tokens\", total=len(df)\n    ):\n        row_fts = fts[row.id]\n\n        inputs = TOKENIZER.encode_plus(\n            row.source,\n            None,\n            add_special_tokens=True,\n            max_length=MD_MAX_LEN,\n            padding=\"max_length\",\n            return_token_type_ids=True,\n            truncation=True,\n        )\n        code_inputs = TOKENIZER.batch_encode_plus(\n            [str(x) for x in row_fts[\"codes\"]] or [\"\"],\n            add_special_tokens=True,\n            max_length=23,\n            padding=\"max_length\",\n            truncation=True,\n        )\n\n        ids = inputs[\"input_ids\"]\n        for x in code_inputs[\"input_ids\"]:\n            ids.extend(x[:-1])\n        ids = ids[:TOTAL_MAX_LEN]\n        if len(ids) != TOTAL_MAX_LEN:\n            ids = ids + [\n                TOKENIZER.pad_token_id,\n            ] * (TOTAL_MAX_LEN - len(ids))\n\n        mask = inputs[\"attention_mask\"]\n        for x in code_inputs[\"attention_mask\"]:\n            mask.extend(x[:-1])\n        mask = mask[:TOTAL_MAX_LEN]\n        if len(mask) != TOTAL_MAX_LEN:\n            mask = mask + [\n                TOKENIZER.pad_token_id,\n            ] * (TOTAL_MAX_LEN - len(mask))\n\n        input_ids[i] = ids\n        attention_mask[i] = mask\n        features[i] = (\n            row_fts[\"total_md\"] / (row_fts[\"total_md\"] + row_fts[\"total_code\"]) or 1\n        )\n\n    return {\n        \"input_ids\": input_ids,\n        \"attention_mask\": attention_mask,\n        \"features\": features,\n    }\n\n\ndef get_ranks(base: pd.Series, derived: List[str]) -> List[str]:\n    return [base.index(d) for d in derived]\n\n\ndef get_dataset(\n    input_ids: np.array,\n    attention_mask: np.array,\n    feature: np.array,\n) -> tf.data.Dataset:\n    dataset = tf.data.Dataset.from_tensor_slices(\n        {\"input_ids\": input_ids, \"attention_mask\": attention_mask, \"feature\": feature}\n    )\n    dataset = dataset.batch(BATCH_SIZE)\n    return dataset.prefetch(tf.data.AUTOTUNE)\n\n\ndef get_model() -> tf.keras.Model:\n    backbone = transformers.TFAutoModel.from_pretrained(BASE_MODEL)\n    input_ids = tf.keras.layers.Input(\n        shape=(TOTAL_MAX_LEN,),\n        dtype=tf.int32,\n        name=\"input_ids\",\n    )\n    attention_mask = tf.keras.layers.Input(\n        shape=(TOTAL_MAX_LEN,),\n        dtype=tf.int32,\n        name=\"attention_mask\",\n    )\n    feature = tf.keras.layers.Input(\n        shape=(1,),\n        dtype=tf.float32,\n        name=\"feature\",\n    )\n    x = backbone({\"input_ids\": input_ids, \"attention_mask\": attention_mask})[0]\n    x = tf.concat([x[:, 0, :], feature], axis=1)\n    outputs = tf.keras.layers.Dense(1, activation=\"linear\", dtype=\"float32\")(x)\n    return tf.keras.Model(\n        inputs=[input_ids, attention_mask, feature],\n        outputs=outputs,\n    )","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-06-22T14:57:52.270955Z","iopub.execute_input":"2022-06-22T14:57:52.271671Z","iopub.status.idle":"2022-06-22T14:57:53.159057Z","shell.execute_reply.started":"2022-06-22T14:57:52.271633Z","shell.execute_reply":"2022-06-22T14:57:53.158106Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Collect Data","metadata":{}},{"cell_type":"code","source":"paths = glob.glob(os.path.join(INPUT_PATH, \"test\", \"*.json\"))\ndf = (\n    pd.concat([read_notebook(x) for x in tqdm(paths, desc=\"Concat\")])\n    .set_index(\"id\", append=True)\n    .swaplevel()\n    .sort_index(level=\"id\", sort_remaining=False)\n).reset_index()\ndf[\"source\"] = df[\"source\"].str.slice(0, MD_MAX_LEN)\ndf[\"rank\"] = df.groupby([\"id\", \"cell_type\"]).cumcount()\ndf[\"pct_rank\"] = df.groupby([\"id\", \"cell_type\"])[\"rank\"].rank(pct=True)\n\nfts = get_features(df)","metadata":{"execution":{"iopub.status.busy":"2022-06-22T14:57:53.16038Z","iopub.execute_input":"2022-06-22T14:57:53.160729Z","iopub.status.idle":"2022-06-22T14:57:53.313812Z","shell.execute_reply.started":"2022-06-22T14:57:53.160695Z","shell.execute_reply":"2022-06-22T14:57:53.312895Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Run Inference","metadata":{}},{"cell_type":"code","source":"with STRATEGY.scope():\n    model = get_model()\n    model.load_weights(\"../input/ai4code-codebert-weights/model_0.h5\")\n\npredict = np.array([], dtype=np.float32)\n\nfor chunk in tqdm(\n    np.array_split(df[df[\"cell_type\"] == \"markdown\"], SLICES), total=SLICES\n):\n    if chunk.empty:\n        continue\n\n    data = tokenize(chunk, fts)\n\n    dataset = get_dataset(data[\"input_ids\"], data[\"attention_mask\"], data[\"features\"])\n    predict = np.r_[\n        predict,\n        model.predict(dataset).reshape(\n            -1,\n        ),\n    ]","metadata":{"execution":{"iopub.status.busy":"2022-06-22T14:57:53.316043Z","iopub.execute_input":"2022-06-22T14:57:53.316899Z","iopub.status.idle":"2022-06-22T14:58:21.72131Z","shell.execute_reply.started":"2022-06-22T14:57:53.316862Z","shell.execute_reply":"2022-06-22T14:58:21.720377Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Save Submission","metadata":{}},{"cell_type":"code","source":"df.loc[df[\"cell_type\"] == \"markdown\", \"pct_rank\"] = predict\ndf = df.sort_values(\"pct_rank\").groupby(\"id\")[\"cell_id\"].apply(\" \".join)\ndf.name = \"cell_order\"\ndf.to_csv(\"submission.csv\")","metadata":{"execution":{"iopub.status.busy":"2022-06-22T14:58:21.722996Z","iopub.execute_input":"2022-06-22T14:58:21.723973Z","iopub.status.idle":"2022-06-22T14:58:21.736328Z","shell.execute_reply.started":"2022-06-22T14:58:21.723933Z","shell.execute_reply":"2022-06-22T14:58:21.73543Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Next Steps\n\nTweak previous steps and beat my score!\n\n* [1/3] [Data Preparation][1] (~5 hours)\n* [2/3] [TPU Training][2] (~4 hours)\n* <span style=\"color:lightgray\">[3/3] GPU Inference ← (you're here)</span>\n\n\n[1]: https://www.kaggle.com/nickuzmenkov/ai4code-tf-tpu-codebert-data-preparation\n[2]: https://www.kaggle.com/nickuzmenkov/ai4code-tf-tpu-codebert-training","metadata":{}}]}