{"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":"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-17T06:41:53.467909Z","iopub.execute_input":"2022-07-17T06:41:53.468222Z","iopub.status.idle":"2022-07-17T06:41:53.781749Z","shell.execute_reply.started":"2022-07-17T06:41:53.468197Z","shell.execute_reply":"2022-07-17T06:41:53.780826Z"},"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-17T06:41:53.783631Z","iopub.execute_input":"2022-07-17T06:41:53.784256Z","iopub.status.idle":"2022-07-17T06:41:53.789529Z","shell.execute_reply.started":"2022-07-17T06:41:53.784229Z","shell.execute_reply":"2022-07-17T06:41:53.788633Z"},"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\nMODEL_NAME = \"distilbert-base-uncased\"\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-17T06:41:53.790724Z","iopub.execute_input":"2022-07-17T06:41:53.792107Z","iopub.status.idle":"2022-07-17T06:41:54.765132Z","shell.execute_reply.started":"2022-07-17T06:41:53.792045Z","shell.execute_reply":"2022-07-17T06:41:54.763986Z"},"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-17T06:41:54.767800Z","iopub.execute_input":"2022-07-17T06:41:54.768140Z","iopub.status.idle":"2022-07-17T06:41:54.792836Z","shell.execute_reply.started":"2022-07-17T06:41:54.768108Z","shell.execute_reply":"2022-07-17T06:41:54.791823Z"},"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-17T06:41:54.793943Z","iopub.execute_input":"2022-07-17T06:41:54.795014Z","iopub.status.idle":"2022-07-17T06:42:05.748407Z","shell.execute_reply.started":"2022-07-17T06:41:54.794979Z","shell.execute_reply":"2022-07-17T06:42:05.747135Z"},"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-17T06:42:05.749668Z","iopub.execute_input":"2022-07-17T06:42:05.749914Z","iopub.status.idle":"2022-07-17T06:42:05.938113Z","shell.execute_reply.started":"2022-07-17T06:42:05.749890Z","shell.execute_reply":"2022-07-17T06:42:05.936431Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = 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    data = tokenize(df.iloc[split], 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\n    for split, index in tqdm(\n        enumerate(np.array_split(np.arange(data[\"labels\"].shape[0]), FILES_PER_FOLD)),\n        desc=f\"Saving\",\n        total=FILES_PER_FOLD,\n    ):\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-17T06:42:05.941385Z","iopub.execute_input":"2022-07-17T06:42:05.941736Z"},"trusted":true},"execution_count":null,"outputs":[]}]}