{"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>[2/3] AI4Code TensorFlow TPU with CodeBert - Training</center></h1>\n\nThis is the second part of my **AI4Code TensorFlow TPU with CodeBert** series:\n\n* [1/3] [Data Preparation][1] (~5 hours)\n* **[2/3] TPU Training ← (you're here)**\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 this version takes only 50 minutes per epoch on Kaggle TPU, so it's lightning fast ⚡.\n\nModel weights are already saved to the dataset **[AI4Code CodeBert Weights][6]**.\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-Tokens][7]**: output from **[Data Preparation][1]** step\n\n### Warning\n\nThis notebook uses Kaggle environment variables. If you run it on Google Colab make sure you explicitly set `VERBOSE` hyperparameter to either 1 or 2.\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-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-weights\n[7]: https://www.kaggle.com/datasets/nickuzmenkov/ai4code-codebert-tokens\n\n# Setup","metadata":{}},{"cell_type":"code","source":"import os\nfrom typing import List\n\nimport numpy as np\nimport pandas as pd\nimport plotly.graph_objects as go\nimport tensorflow as tf\nimport transformers\nfrom kaggle_datasets import KaggleDatasets\nfrom sklearn.model_selection import KFold","metadata":{"execution":{"iopub.status.busy":"2022-07-20T06:06:22.242251Z","iopub.execute_input":"2022-07-20T06:06:22.242889Z","iopub.status.idle":"2022-07-20T06:06:30.627496Z","shell.execute_reply.started":"2022-07-20T06:06:22.242796Z","shell.execute_reply":"2022-07-20T06:06:30.626476Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"RANDOM_STATE = 42\nN_SPLITS = 5\nTOTAL_MAX_LEN = 512\nBASE_MODEL = \"microsoft/codebert-base\"\nGCS_PATH = KaggleDatasets().get_gcs_path(\"ai4code-codebert-tokens\")\nEPOCHS = 10\nLR = 3e-5\nWARMUP_RATE = 0.05\nVERBOSE = 1 if os.environ[\"KAGGLE_KERNEL_RUN_TYPE\"] == \"Interactive\" else 2\n\ntry:\n    TPU = tf.distribute.cluster_resolver.TPUClusterResolver()\n    tf.config.experimental_connect_to_cluster(TPU)\n    tf.tpu.experimental.initialize_tpu_system(TPU)\n    STRATEGY = tf.distribute.experimental.TPUStrategy(TPU)\n    BATCH_SIZE = 64 * STRATEGY.num_replicas_in_sync\nexcept Exception:\n    TPU = None\n    STRATEGY = tf.distribute.get_strategy()\n    BATCH_SIZE = 4\n\nprint(\"TensorFlow\", tf.__version__)\n\nif TPU is not None:\n    print(\"Using TPU v3-8\")\nelse:\n    print(\"Using GPU/CPU\")\n\nprint(\"Batch size:\", BATCH_SIZE)","metadata":{"execution":{"iopub.status.busy":"2022-07-20T06:06:30.630562Z","iopub.execute_input":"2022-07-20T06:06:30.631287Z","iopub.status.idle":"2022-07-20T06:06:36.949313Z","shell.execute_reply.started":"2022-07-20T06:06:30.631235Z","shell.execute_reply":"2022-07-20T06:06:36.948599Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def count_samples(filenames: List[str]) -> int:\n    return sum(int(os.path.basename(x).split(\".\")[0].split(\"-\")[-1]) for x in filenames)\n\n\ndef read_tfrecord(example: tf.Tensor) -> tf.Tensor:\n    features = {\n        \"input_ids\": tf.io.FixedLenFeature([TOTAL_MAX_LEN], tf.int64),\n        \"attention_mask\": tf.io.FixedLenFeature([TOTAL_MAX_LEN], tf.int64),\n        \"feature\": tf.io.FixedLenFeature([], tf.float32),\n        \"label\": tf.io.FixedLenFeature([], tf.float32),\n    }\n    example = tf.io.parse_single_example(example, features)\n    return (\n        {\n            \"input_ids\": tf.cast(example[\"input_ids\"], tf.int32),\n            \"attention_mask\": tf.cast(example[\"attention_mask\"], tf.int32),\n            \"feature\": example[\"feature\"],\n        },\n        example[\"label\"],\n    )\n\n\ndef get_dataset(\n    filenames: List[str],\n    ordered: bool = False,\n    repeated: bool = True,\n    cached: bool = False,\n) -> tf.data.Dataset:\n    auto = tf.data.experimental.AUTOTUNE\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=auto)\n    if not ordered:\n        ignore_order = tf.data.Options()\n        ignore_order.experimental_deterministic = False\n        dataset = dataset.with_options(ignore_order)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=auto)\n    if not ordered:\n        dataset = dataset.shuffle(2048, seed=RANDOM_STATE)\n    if repeated:\n        dataset = dataset.repeat()\n    dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)\n    if cached:\n        dataset = dataset.cache()\n    dataset = dataset.prefetch(auto)\n    return STRATEGY.experimental_distribute_dataset(dataset)\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    )\n\n\nclass WarmupLinearDecay(tf.keras.optimizers.schedules.LearningRateSchedule):\n    def __init__(\n        self,\n        base_learning_rate: float,\n        warmup_steps: int,\n        total_steps: int,\n    ) -> None:\n        self._base_learning_rate = base_learning_rate\n        self._warmup_steps = warmup_steps\n        self._total_steps = total_steps\n\n    def __call__(self, step: int) -> float:\n        return self._base_learning_rate * tf.cond(\n            tf.math.less_equal(step, warmup_steps),\n            lambda: step / self._warmup_steps,\n            lambda: (step - total_steps) / (self._warmup_steps - self._total_steps),\n        )","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-07-20T06:06:36.951002Z","iopub.execute_input":"2022-07-20T06:06:36.951323Z","iopub.status.idle":"2022-07-20T06:06:36.984316Z","shell.execute_reply.started":"2022-07-20T06:06:36.951283Z","shell.execute_reply":"2022-07-20T06:06:36.983025Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"for i, (train_index, val_index) in enumerate(KFold(n_splits=N_SPLITS).split(range(N_SPLITS))):\n    if TPU is not None:\n        tf.tpu.experimental.initialize_tpu_system(TPU)\n\n    train_filenames = np.ravel(\n        [\n            tf.io.gfile.glob(os.path.join(GCS_PATH, \"tfrec\", str(x), \"*.tfrec\"))\n            for x in train_index\n        ]\n    )\n    steps_per_epoch = count_samples(train_filenames) // BATCH_SIZE\n    train_dataset = get_dataset(train_filenames)\n\n    val_filenames = np.ravel(\n        [\n            tf.io.gfile.glob(os.path.join(GCS_PATH, \"tfrec\", str(x), \"*.tfrec\"))\n            for x in val_index\n        ]\n    )\n    validation_steps = count_samples(val_filenames) // BATCH_SIZE\n    val_dataset = get_dataset(val_filenames, ordered=True, repeated=False, cached=True)\n\n    with STRATEGY.scope():\n        model = get_model()\n\n        total_steps = steps_per_epoch * EPOCHS\n        warmup_steps = int(WARMUP_RATE * total_steps)\n\n        optimizer = transformers.AdamWeightDecay(\n            learning_rate=WarmupLinearDecay(\n                base_learning_rate=LR,\n                warmup_steps=warmup_steps,\n                total_steps=total_steps,\n            ),\n            weight_decay_rate=0.01,\n            exclude_from_weight_decay=[\n                \"bias\",\n                \"LayerNorm.bias\",\n                \"LayerNorm.weight\",\n            ],\n        )\n        model.compile(loss=\"mae\", optimizer=optimizer)\n\n    metrics = model.fit(\n        train_dataset,\n        steps_per_epoch=steps_per_epoch,\n        validation_data=val_dataset,\n        validation_steps=validation_steps,\n        epochs=EPOCHS,\n        verbose=VERBOSE,\n    ).history\n\n    model.save_weights(f\"model_{i}.h5\")\n    break","metadata":{"execution":{"iopub.status.busy":"2022-07-20T06:06:36.986431Z","iopub.execute_input":"2022-07-20T06:06:36.987009Z","iopub.status.idle":"2022-07-20T06:09:00.152281Z","shell.execute_reply.started":"2022-07-20T06:06:36.986962Z","shell.execute_reply":"2022-07-20T06:09:00.150992Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Metrics","metadata":{}},{"cell_type":"code","source":"metrics = pd.DataFrame(metrics)\n\ngo.Figure(\n    data=(\n        go.Scatter(x=metrics.index, y=metrics[\"loss\"], name=\"train\"),\n        go.Scatter(x=metrics.index, y=metrics[\"val_loss\"], name=\"validation\"),\n    ),\n    layout=dict(\n        width=600,\n        title_text=\"Model loss\",\n        xaxis_title_text=\"Epoch\",\n        font=dict(size=16),\n    ),\n)","metadata":{"execution":{"iopub.status.busy":"2022-07-20T06:10:27.569741Z","iopub.execute_input":"2022-07-20T06:10:27.570351Z","iopub.status.idle":"2022-07-20T06:10:27.594115Z","shell.execute_reply.started":"2022-07-20T06:10:27.570314Z","shell.execute_reply":"2022-07-20T06:10:27.593037Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Next Steps\n\nGo to the model weights dataset **[here][1]** or continue exploring:\n\n* [1/3] [Data Preparation][2] (~3 hours)\n* <span style=\"color:lightgray\">[2/3] TPU Training ← (you're here)</span>\n* [3/3] [GPU Inference][3] (~2 hours)\n\n\n[1]: https://www.kaggle.com/datasets/nickuzmenkov/ai4code-codebert-weights\n[2]: https://www.kaggle.com/nickuzmenkov/ai4code-tf-tpu-codebert-data-preparation\n[3]: https://www.kaggle.com/nickuzmenkov/ai4code-tf-tpu-codebert-inference","metadata":{}}]}