{"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>[1/3] AI4Code TensorFlow TPU with CodeBert - Data Preparation</center></h1>\n\nThis is the first part of my **AI4Code TensorFlow TPU with CodeBert** series:\n\n* **[1/3] Data Preparation ← (you're here)**\n* [2/3] [TPU Training][1] (~4 hours)\n* [3/3] [GPU Inference][2] (~2 hours)\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][1]** takes only 50 minutes per epoch on Kaggle TPU, so it's lightning fast ⚡.\n\nOutputs of this notebook are already saved to the dataset **[AI4Code CodeBert Tokens][6]** so feel free to skip this part unless you need to customize it!\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### Warning\n\nThis notebook uses all data only when submitted via **Save & Run All (Commit)** and only the first 1,000 notebook entries in interactive session. This behaviour is bound to Kaggle environment variables. To make it process all data on Google Colab or your local machine, please explicitly set the `LIMIT` hyperparameter to `None`.\n\n[1]: https://www.kaggle.com/nickuzmenkov/ai4code-tf-tpu-codebert-training\n[2]: https://www.kaggle.com/nickuzmenkov/ai4code-tf-tpu-codebert-inference\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-tokens\n\n# Setup","metadata":{"_uuid":"b4acf22f-a811-4675-aa72-f5142f65d072","_cell_guid":"c3ef605e-93ea-48da-87fd-5f0d38246e49","trusted":true}},{"cell_type":"code","source":"!mkdir 'raw' 'tfrec'","metadata":{"_uuid":"35729e97-72d2-4239-a596-560616f36d65","_cell_guid":"487f5be4-716f-400c-9d45-cd3df27d8e3f","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-07-23T10:05:08.139451Z","iopub.execute_input":"2022-07-23T10:05:08.140169Z","iopub.status.idle":"2022-07-23T10:05:08.934396Z","shell.execute_reply.started":"2022-07-23T10:05:08.140075Z","shell.execute_reply":"2022-07-23T10:05:08.933033Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import glob\nimport json\nimport os\nfrom typing import List\n\nimport numpy as np\nimport pandas as pd\nimport tensorflow as tf\nimport transformers\nfrom sklearn.model_selection import GroupKFold\nfrom sklearn.utils import shuffle\nfrom tqdm.notebook import tqdm","metadata":{"_uuid":"54b466e4-e2a1-40a1-b745-2c8eeb3e9455","_cell_guid":"ea09849c-12e7-4045-a497-eae98ced6f1e","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-07-23T10:05:08.937173Z","iopub.execute_input":"2022-07-23T10:05:08.937549Z","iopub.status.idle":"2022-07-23T10:05:20.795034Z","shell.execute_reply.started":"2022-07-23T10:05:08.937511Z","shell.execute_reply":"2022-07-23T10:05:20.793947Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"RANDOM_STATE = 42\nMD_MAX_LEN = 64\nTOTAL_MAX_LEN = 512\nK_FOLDS = 5\nFILES_PER_FOLD = 16\nLIMIT = 1_000 if os.environ[\"KAGGLE_KERNEL_RUN_TYPE\"] == \"Interactive\" else None\n# MODEL_NAME = \"microsoft/codebert-base\"\nMODEL_NAME = \"bert-base-multilingual-cased\"\nTOKENIZER = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)\nINPUT_PATH = \"../input/AI4Code\"","metadata":{"_uuid":"7a407cdd-5a12-46fa-8c25-db31ed69cd55","_cell_guid":"84bb4982-7ae3-400b-8288-212390c9c0b4","collapsed":false,"_kg_hide-output":true,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-07-23T10:05:20.796528Z","iopub.execute_input":"2022-07-23T10:05:20.797237Z","iopub.status.idle":"2022-07-23T10:05:36.904884Z","shell.execute_reply.started":"2022-07-23T10:05:20.797201Z","shell.execute_reply":"2022-07-23T10:05:36.903657Z"},"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    labels = 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        labels[i] = row.pct_rank\n\n    return {\n        \"input_ids\": input_ids,\n        \"attention_mask\": attention_mask,\n        \"features\": features,\n        \"labels\": labels,\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 _serialize_sample(\n    input_ids: np.array,\n    attention_mask: np.array,\n    feature: np.float64,\n    label: np.float64,\n) -> bytes:\n    feature = {\n        \"input_ids\": tf.train.Feature(int64_list=tf.train.Int64List(value=input_ids)),\n        \"attention_mask\": tf.train.Feature(\n            int64_list=tf.train.Int64List(value=attention_mask)\n        ),\n        \"feature\": tf.train.Feature(float_list=tf.train.FloatList(value=[feature])),\n        \"label\": tf.train.Feature(float_list=tf.train.FloatList(value=[label])),\n    }\n    sample = tf.train.Example(features=tf.train.Features(feature=feature))\n    return sample.SerializeToString()\n\n\ndef serialize(\n    input_ids: np.array,\n    attention_mask: np.array,\n    features: np.array,\n    labels: np.array,\n    path: str,\n) -> None:\n    with tf.io.TFRecordWriter(path) as writer:\n        for args in zip(input_ids, attention_mask, features, labels):\n            writer.write(_serialize_sample(*args))","metadata":{"_uuid":"54183c98-306a-4b06-89f8-627ac7fc3cec","_cell_guid":"b47c2f25-8906-4e85-a8e1-b5fcd8c94703","collapsed":false,"_kg_hide-input":true,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-07-23T10:05:36.907736Z","iopub.execute_input":"2022-07-23T10:05:36.908376Z","iopub.status.idle":"2022-07-23T10:05:36.943670Z","shell.execute_reply.started":"2022-07-23T10:05:36.908334Z","shell.execute_reply":"2022-07-23T10:05:36.942308Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Collect Data","metadata":{"_uuid":"1ecc70eb-0f05-475f-b54e-724686873288","_cell_guid":"91005f15-8d97-47ca-9887-889c22677bc9","trusted":true}},{"cell_type":"code","source":"paths = glob.glob(os.path.join(INPUT_PATH, \"train\", \"*.json\"))\nif LIMIT is not None:\n    paths = paths[:LIMIT]\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)\n\ndf_orders = pd.read_csv(\n    os.path.join(INPUT_PATH, \"train_orders.csv\"),\n    index_col=\"id\",\n    squeeze=True,\n).str.split()\ndf_orders_ = df_orders.to_frame().join(\n    df.reset_index(\"cell_id\").groupby(\"id\")[\"cell_id\"].apply(list),\n    how=\"right\",\n)\n\nranks = {}\nfor id_, cell_order, cell_id in df_orders_.itertuples():\n    ranks[id_] = {\"cell_id\": cell_id, \"rank\": get_ranks(cell_order, cell_id)}\ndf_ranks = (\n    pd.DataFrame.from_dict(ranks, orient=\"index\")\n    .rename_axis(\"id\")\n    .apply(pd.Series.explode)\n    .set_index(\"cell_id\", append=True)\n)\n\ndf_ancestors = pd.read_csv(\n    os.path.join(INPUT_PATH, \"train_ancestors.csv\"), index_col=\"id\"\n)\ndf = (\n    df.reset_index()\n    .merge(df_ranks, on=[\"id\", \"cell_id\"])\n    .merge(df_ancestors, on=[\"id\"])\n)\n\ndf[\"pct_rank\"] = df[\"rank\"] / df.groupby(\"id\")[\"cell_id\"].transform(\"count\")\ndf = df.sort_values(\"pct_rank\").reset_index(drop=True)\n\nfeatures = get_features(df)\n\ndf = df[df[\"cell_type\"] == \"markdown\"]\ndf = df.drop([\"rank\", \"parent_id\", \"cell_type\"], axis=1).dropna()","metadata":{"_uuid":"e5bbb7e7-2ed2-4595-9398-53d997cfbba9","_cell_guid":"51f46e7b-b4e8-400d-944b-954bcc5685f2","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-07-23T10:05:36.945426Z","iopub.execute_input":"2022-07-23T10:05:36.945904Z","iopub.status.idle":"2022-07-23T10:05:55.903591Z","shell.execute_reply.started":"2022-07-23T10:05:36.945848Z","shell.execute_reply":"2022-07-23T10:05:55.902405Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Make Tokens & Save","metadata":{"_uuid":"99ede939-5cf8-4bd9-8762-a0712dd91cc4","_cell_guid":"087936e6-ad8e-49d4-a619-b61289dc6a48","trusted":true}},{"cell_type":"code","source":"df.to_csv(\"data.csv\")\nwith open(\"features.json\", \"w\") as file:\n    json.dump(features, file)","metadata":{"_uuid":"3fe34d41-47f7-4449-be94-f1ab4faefc85","_cell_guid":"5dcfae2f-7131-4918-987f-49182b3d9157","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-07-23T10:05:55.905158Z","iopub.execute_input":"2022-07-23T10:05:55.905913Z","iopub.status.idle":"2022-07-23T10:05:56.253760Z","shell.execute_reply.started":"2022-07-23T10:05:55.905856Z","shell.execute_reply":"2022-07-23T10:05:56.252939Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import json\n\ndf = shuffle(df, random_state=RANDOM_STATE)\n\nfor fold, (_, split) in enumerate(\n    GroupKFold(K_FOLDS).split(df, groups=df[\"ancestor_id\"])\n):\n    print(\"=\" * 36, f\"Fold {fold}\", \"=\" * 36)\n    fold_dir = f\"tfrec/{fold}\"\n    if not os.path.exists(fold_dir):\n        os.mkdir(fold_dir)\n\n    df_sub = df.iloc[split]\n#     data = tokenize(df_sub, features)\n\n#     np.savez_compressed(\n#         f\"raw/{fold}.npz\",\n#         input_ids=data[\"input_ids\"],\n#         attention_mask=data[\"attention_mask\"],\n#         features=data[\"features\"],\n#         labels=data[\"labels\"],\n#     )\n#     print(len(df_sub), data[\"labels\"].shape[0])\n\n    for split, index in tqdm(\n        enumerate(np.array_split(np.arange(len(df_sub)), FILES_PER_FOLD)),\n        desc=f\"Saving\",\n        total=FILES_PER_FOLD,\n    ):\n        ids = df_sub.iloc[index][\"id\"].unique()\n        print(len(ids))\n        ids_file = os.path.join(fold_dir, f\"{split:02d}-{len(index):06d}.json\")\n        with open(ids_file, \"w\") as f:\n            json.dump(list(ids), f, indent=4)\n#         serialize(\n#             input_ids=data[\"input_ids\"][index],\n#             attention_mask=data[\"attention_mask\"][index],\n#             features=data[\"features\"][index],\n#             labels=data[\"labels\"][index],\n#             path=os.path.join(fold_dir, f\"{split:02d}-{len(index):06d}.tfrec\"),\n#         )","metadata":{"_uuid":"d0775a9c-bd53-4286-868c-4e9fa13b57b7","_cell_guid":"3962a968-4330-4ed7-9aaa-0116a3d79def","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-07-23T10:13:40.005357Z","iopub.execute_input":"2022-07-23T10:13:40.005722Z","iopub.status.idle":"2022-07-23T10:13:40.284452Z","shell.execute_reply.started":"2022-07-23T10:13:40.005692Z","shell.execute_reply":"2022-07-23T10:13:40.283641Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Next Steps\n\nGo to the results dataset **[AI4Code CodeBert Tokens][3]** or continue exploring:\n\n* <span style=\"color:lightgray\">[1/3] Data Preparation ← (you're here)</span>\n* [2/3] [TPU Training][1] (~4 hours)\n* [3/3] [GPU Inference][2] (~2 hours)\n\n\n[1]: https://www.kaggle.com/nickuzmenkov/ai4code-tf-tpu-codebert-training\n[2]: https://www.kaggle.com/nickuzmenkov/ai4code-tf-tpu-codebert-inference\n[3]: https://www.kaggle.com/datasets/nickuzmenkov/ai4code-codebert-tokens","metadata":{"_uuid":"e7290989-4a74-4be6-b6f0-22d566076b79","_cell_guid":"07eab696-a6aa-4426-9057-40c750d0bb01","trusted":true}}]}