{"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":"# Intro\n\nThis notebook is basically a translation of __[AI4Code Pytorch DistilBert Baseline][0]__ to help out TensorFlow users. I shortened the original work leaving only its bare bones - __[Data Preparation][1]__, __[Model Training][2]__, and __[Inference][3]__.\n\n### TPU usage\n\nThis notebook is updated for TPU usage: just select TPU v3-8 as accelerator to train your model with all the competition data and get 5x performance boost! \n\n__Warning__: accroding to the competition rules you won't be able to submit your notebook if it uses TPU. Train your model in a separate notebook as a workaround.\n\n# Setup\n\n[0]: https://www.kaggle.com/code/aerdem4/ai4code-pytorch-distilbert-baseline\n[1]: #data_preparation\n[2]: #training\n[3]: #inference","metadata":{"_uuid":"c34597d0-7dc4-4310-ad13-8aa28d2b780c","_cell_guid":"0bd99a87-8ce9-48f3-9efd-bc84eba430c5","trusted":true}},{"cell_type":"code","source":"import glob\nimport json\nimport os\nfrom typing import Optional, Tuple\n\nimport numpy as np\nimport pandas as pd\nimport tensorflow as tf\nimport transformers\nfrom IPython.display import display\nfrom sklearn.utils import shuffle\nfrom sklearn.model_selection import GroupKFold\nfrom tqdm.notebook import tqdm","metadata":{"_uuid":"6e333246-7066-4549-896b-e28104f40551","_cell_guid":"b0d814ae-97f7-4e28-9cb5-8ab805c32301","collapsed":false,"_kg_hide-output":true,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DATA_PATH = \"../input/AI4Code\"\nBASE_MODEL = \"../input/huggingface-bert-variants/distilbert-base-uncased/distilbert-base-uncased\"\nN_SPLITS = 5\nSEQ_LEN = 128\nRANDOM_STATE = 42\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 = 128 * STRATEGY.num_replicas_in_sync\nexcept Exception:\n    TPU = None\n    STRATEGY = tf.distribute.get_strategy()\n    BATCH_SIZE = 32\n    LIMIT = 10_000\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":{"_uuid":"109544fc-035f-45fc-bb2a-44546597915a","_cell_guid":"f5cc19ad-8fc8-4f5d-8d43-684c739f47ca","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def read_notebook(path: str) -> pd.DataFrame:\n    with open(path) as file:\n        df = pd.DataFrame(json.load(file))\n    df[\"id\"] = os.path.splitext(os.path.basename(path))[0]\n    return df\n\n\ndef expand_order(row: Tuple[str, str]) -> pd.DataFrame:\n    cell_ids = row[1].split(\" \")\n    df = pd.DataFrame(\n        {\n            \"id\": [row[0] for _ in range(len(cell_ids))],\n            \"cell_id\": cell_ids,\n            \"rank\": range(len(cell_ids)),\n        }\n    )\n    df[\"pct_rank\"] = df[\"rank\"] / len(df)\n    return df\n\n\ndef tokenize(source: pd.Series) -> Tuple[np.array, np.array]:\n    tokenizer = transformers.AutoTokenizer.from_pretrained(BASE_MODEL, do_lower_case=True)\n\n    input_ids = np.zeros((len(source), SEQ_LEN), dtype=\"int32\")\n    attention_mask = np.zeros((len(source), SEQ_LEN), dtype=\"int32\")\n\n    for i, x in enumerate(tqdm(source, total=len(source))):\n        encoding = tokenizer.encode_plus(\n            x,\n            None,\n            add_special_tokens=True,\n            max_length=SEQ_LEN,\n            padding=\"max_length\",\n            return_token_type_ids=True,\n            truncation=True,\n        )\n        input_ids[i] = encoding[\"input_ids\"]\n        attention_mask[i] = encoding[\"attention_mask\"]\n\n    return input_ids, attention_mask\n\n\ndef get_dataset(\n    input_ids: np.array,\n    attention_mask: np.array,\n    labels: Optional[np.array] = None,\n    ordered: bool = False,\n    repeated: bool = False,\n) -> tf.data.Dataset:\n    if labels is not None:\n        dataset = tf.data.Dataset.from_tensor_slices(\n            ({\"input_ids\": input_ids, \"attention_mask\": attention_mask}, labels)\n        )\n    else:\n        dataset = tf.data.Dataset.from_tensor_slices(\n            {\"input_ids\": input_ids, \"attention_mask\": attention_mask}\n        )\n    if repeated:\n        dataset = dataset.repeat()\n    if not ordered:\n        dataset = dataset.shuffle(1024)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.prefetch(tf.data.AUTOTUNE)\n    return dataset\n\n\ndef get_model() -> tf.keras.Model:\n    backbone = transformers.TFDistilBertModel.from_pretrained(BASE_MODEL)\n    input_ids = tf.keras.layers.Input(\n        shape=(SEQ_LEN,),\n        dtype=tf.int32,\n        name=\"input_ids\",\n    )\n    attention_mask = tf.keras.layers.Input(\n        shape=(SEQ_LEN,),\n        dtype=tf.int32,\n        name=\"attention_mask\",\n    )\n    x = backbone(\n        {\n            \"input_ids\": input_ids,\n            \"attention_mask\": attention_mask,\n        },\n    )\n    outputs = tf.keras.layers.Dense(1, activation=\"linear\", dtype=\"float32\")(x[0][:, 0, :])\n\n    model = tf.keras.Model(\n        inputs=[input_ids, attention_mask],\n        outputs=outputs,\n    )\n    model.compile(\n        optimizer=tf.keras.optimizers.Adam(learning_rate=5e-5),\n        loss=tf.keras.losses.MeanSquaredError(),\n    )\n    return model","metadata":{"_uuid":"ef219d83-591d-4b8e-a859-6bf48bcdf868","_cell_guid":"1be1dec6-5a4e-41bb-a59c-409bdfaa9d22","collapsed":false,"_kg_hide-input":true,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# <a id=\"data_preparation\">Data Preparation</a>\n\nThe idea behind this is simple:\n    \n1. Get only markdown cells from the training notebooks\n2. Teach the model to predict relative position (or `pct_rank`, or normalized rank) of the markdown cell given only its text and no other context","metadata":{"_uuid":"2645e1f8-24ea-4674-83d9-1e6efd69baa0","_cell_guid":"cf6f2366-bf39-4116-b0e7-4c695ef6bd61","trusted":true}},{"cell_type":"code","source":"paths = glob.glob(os.path.join(DATA_PATH, \"train\", \"*.json\"))\nif LIMIT is not None:\n    paths = paths[:LIMIT]\n\nsource_df = pd.concat([read_notebook(x) for x in tqdm(paths, total=len(paths))])\n\nsource_df = source_df[source_df[\"cell_type\"] == \"markdown\"]\nsource_df = source_df.drop(\"cell_type\", axis=1)\nsource_df = source_df.rename_axis(\"cell_id\").reset_index()\n\norder_df = pd.read_csv(os.path.join(DATA_PATH, \"train_orders.csv\"), index_col=\"id\")\norder_df = pd.concat(\n    [expand_order(row) for row in tqdm(order_df.itertuples(), total=len(order_df))]\n)\n\nancestors_df = pd.read_csv(\n    os.path.join(DATA_PATH, \"train_ancestors.csv\"),\n    usecols=[\"id\", \"ancestor_id\"],\n    index_col=\"id\",\n)\n\ndf = source_df.merge(order_df, on=[\"id\", \"cell_id\"]).merge(ancestors_df, on=\"id\")\ndf = df.dropna()\ndisplay(df)","metadata":{"_uuid":"14b178af-ba80-4c1f-a50b-68e83500d915","_cell_guid":"578b3215-f2ad-446e-a969-096c1cfb8bc9","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"input_ids, attention_mask = tokenize(df[\"source\"])\n\nlabels = df[\"pct_rank\"].to_numpy()\ngroups = df[\"ancestor_id\"].to_numpy()\n\nprint(\"input_ids:\", input_ids.shape)\nprint(\"attention_mask:\", attention_mask.shape)\nprint(\"labels:\", labels.shape)\nprint(\"groups:\", groups.shape)","metadata":{"_uuid":"728ba8c7-3a49-4295-a734-faf2dde4db0b","_cell_guid":"7db26726-14b5-4874-92da-5b95c52df224","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# <a id=\"training\">Model Training</a>\n\nKFold ensembles can get you an extra boost. Just remove the `break` statement at the bottom to run all `N_SPLITS`.","metadata":{"_uuid":"2afb313a-e43d-4008-b86f-e5b56cb2f4a4","_cell_guid":"f869e307-2be8-4688-a7e5-0e38ff8d3c1e","trusted":true}},{"cell_type":"code","source":"input_ids, attention_mask, labels, groups = shuffle(\n    input_ids, attention_mask, labels, groups, random_state=RANDOM_STATE\n)\nkfold = GroupKFold(n_splits=N_SPLITS)\n\nfor i, (train_index, val_index) in enumerate(kfold.split(input_ids, labels, groups=groups)):\n    if TPU is not None:\n        tf.tpu.experimental.initialize_tpu_system(TPU)\n\n    with STRATEGY.scope():\n        model = get_model()\n        model.summary()\n\n    train_dataset = get_dataset(\n        input_ids=input_ids[train_index],\n        attention_mask=attention_mask[train_index],\n        labels=labels[train_index],\n        repeated=True,\n    )\n    val_dataset = get_dataset(\n        input_ids=input_ids[val_index],\n        attention_mask=attention_mask[val_index],\n        labels=labels[val_index],\n        ordered=True,\n    )\n\n    model.fit(\n        train_dataset,\n        validation_data=val_dataset,\n        steps_per_epoch=len(train_index) // BATCH_SIZE,\n        epochs=1,\n        verbose=2,\n    )\n\n    model.save_weights(f\"model_{i}.h5\")\n    break","metadata":{"_uuid":"eb9ff631-639a-4c87-a525-8e3b4b8088de","_cell_guid":"46ce375f-ce87-4b33-8291-b958e0a050ee","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# <a id=\"inference\">Inference</a>\n\nMove this part to a separate notebook to be able to train your models on TPU and load them here simply as:\n\n```python\nmodel = get_model()\nmodel.load_weights(\"../path/to/your/dataset/model_0.h5\")\n```\n\nThere are only 4 notebooks in the test folder available at the runtime, but there will be over 20,000 when you submit it!","metadata":{"_uuid":"440bdd70-8063-40a1-bca5-29367062a625","_cell_guid":"95593826-8940-43e2-a4e4-b67b4ea4adfb","trusted":true}},{"cell_type":"code","source":"paths = glob.glob(os.path.join(DATA_PATH, \"test\", \"*.json\"))\n\ndf = pd.concat([read_notebook(x) for x in tqdm(paths, total=len(paths))])\ndf = df.rename_axis(\"cell_id\").reset_index()\n\ndf[\"rank\"] = df.groupby([\"id\", \"cell_type\"]).cumcount()\ndf[\"pct_rank\"] = df.groupby([\"id\", \"cell_type\"])[\"rank\"].rank(pct=True)\n\ndisplay(df)","metadata":{"_uuid":"c6522cc2-f859-4edf-85e8-c7bea7a650d9","_cell_guid":"26c47186-ba38-4c78-93d3-8737b629c82c","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"input_ids, attention_mask = tokenize(df[df[\"cell_type\"] == \"markdown\"][\"source\"])\ntest_dataset = get_dataset(\n    input_ids=input_ids,\n    attention_mask=attention_mask,\n    ordered=True,\n)\ny_pred = model.predict(test_dataset)","metadata":{"_uuid":"b5f672e2-e8ec-4402-a7f8-46b1647cd93f","_cell_guid":"29a7f191-5ceb-43fb-b450-d0e13ef720bb","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df.loc[df[\"cell_type\"] == \"markdown\", \"pct_rank\"] = y_pred\ndf = df.sort_values(\"pct_rank\").groupby(\"id\", as_index=False)[\"cell_id\"].apply(lambda x: \" \".join(x))\ndf.rename(columns={\"cell_id\": \"cell_order\"}, inplace=True)\ndf.to_csv(\"submission.csv\", index=False)\ndisplay(df)","metadata":{"_uuid":"8ed5cefd-9f92-4641-870e-647a6a4bbf95","_cell_guid":"2679346a-f54a-4776-9633-67f926c33202","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]}]}