{"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":"# Task:\nPredict correct positions of markdown cells in the Jupyter notebook based on the given ordred code cells.\n\n# Input:\n- markdown cells text\n    - unordered\n- code cells text\n    - must be ordered\n\n# Output:\nThe original positions of the markdown cells in the Jupyter notebook","metadata":{}},{"cell_type":"markdown","source":"# Import and Load pipeline","metadata":{}},{"cell_type":"code","source":"import pathlib\nMODEL_ROOT = pathlib.Path(\"/kaggle/input/ai4code-parallel-bert/pytorch/deberta-v3-large/2\")\n!ls $MODEL_ROOT","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:14.136165Z","iopub.execute_input":"2023-07-11T19:31:14.136851Z","iopub.status.idle":"2023-07-11T19:31:15.090667Z","shell.execute_reply.started":"2023-07-11T19:31:14.136817Z","shell.execute_reply":"2023-07-11T19:31:15.089437Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n\nimport sys\nsys.path.append(str(MODEL_ROOT))\nimport ai4code_parallel_bert\nsys.path.pop(-1)","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:16.868492Z","iopub.execute_input":"2023-07-11T19:31:16.868896Z","iopub.status.idle":"2023-07-11T19:31:18.956296Z","shell.execute_reply.started":"2023-07-11T19:31:16.868862Z","shell.execute_reply":"2023-07-11T19:31:18.955355Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"CHECKPOINT_PATH = MODEL_ROOT / \"checkpoint/deberta-v3-large\"\nDEVICE = \"cuda:0\"\nBERT_PROCESS_BATCH_SIZE = 4\n\npipe = ai4code_parallel_bert.pipeline.Pipeline(CHECKPOINT_PATH, DEVICE, BERT_PROCESS_BATCH_SIZE)","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-07-11T19:31:19.825586Z","iopub.execute_input":"2023-07-11T19:31:19.826184Z","iopub.status.idle":"2023-07-11T19:31:36.316368Z","shell.execute_reply.started":"2023-07-11T19:31:19.826145Z","shell.execute_reply":"2023-07-11T19:31:36.315357Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load competition data","metadata":{}},{"cell_type":"code","source":"import pathlib\nCOMPETITION_ROOT = pathlib.Path(\"/kaggle/input/AI4Code/\")\n!ls $COMPETITION_ROOT","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-07-11T19:31:36.318567Z","iopub.execute_input":"2023-07-11T19:31:36.318943Z","iopub.status.idle":"2023-07-11T19:31:37.327001Z","shell.execute_reply.started":"2023-07-11T19:31:36.318909Z","shell.execute_reply":"2023-07-11T19:31:37.325780Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"SAMPLE_TEST_FILE = COMPETITION_ROOT / \"test/0009d135ece78d.json\"\nnotebook_data = ai4code_parallel_bert.pipeline.load_ai4code_json_file(SAMPLE_TEST_FILE)\nnotebook_data.keys()","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:37.328910Z","iopub.execute_input":"2023-07-11T19:31:37.329680Z","iopub.status.idle":"2023-07-11T19:31:37.354574Z","shell.execute_reply.started":"2023-07-11T19:31:37.329640Z","shell.execute_reply":"2023-07-11T19:31:37.353549Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\" ========== first code cell text ==========\")\nprint(notebook_data[\"code\"][\"texts\"][0])\nprint()\n\nprint(\" ========== second code cell text ==========\")\nprint(notebook_data[\"code\"][\"texts\"][1])\nprint()\n\nprint(\" ========== code cell ids ==========\")\nprint(notebook_data[\"code\"][\"cell_ids\"])","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:37.357323Z","iopub.execute_input":"2023-07-11T19:31:37.357742Z","iopub.status.idle":"2023-07-11T19:31:37.365052Z","shell.execute_reply.started":"2023-07-11T19:31:37.357711Z","shell.execute_reply":"2023-07-11T19:31:37.363438Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\" ========== a markdown cell text ==========\")\nprint(notebook_data[\"markdown\"][\"texts\"][0])\nprint()\n\nprint(\" ========== another markdown cell text ==========\")\nprint(notebook_data[\"markdown\"][\"texts\"][1])\nprint()\n\nprint(\" ========== markdown cell ids ==========\")\nprint(notebook_data[\"markdown\"][\"cell_ids\"])","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:37.367340Z","iopub.execute_input":"2023-07-11T19:31:37.368113Z","iopub.status.idle":"2023-07-11T19:31:37.376782Z","shell.execute_reply.started":"2023-07-11T19:31:37.368080Z","shell.execute_reply":"2023-07-11T19:31:37.375870Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Predict","metadata":{}},{"cell_type":"code","source":"pipeline_input = {\n    \"code\": {\n        \"texts\": notebook_data[\"code\"][\"texts\"],\n        \"cell_ids\": notebook_data[\"code\"][\"cell_ids\"],\n    },\n    \"markdown\": {\n        \"texts\": notebook_data[\"markdown\"][\"texts\"],\n        \"cell_ids\": notebook_data[\"markdown\"][\"cell_ids\"],\n    },\n}\n\nimport torch\nwith torch.no_grad():\n  out = pipe(pipeline_input)\n\nprint(\" ========== keys of output dict ==========\")\nprint(out.keys())","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:37.378289Z","iopub.execute_input":"2023-07-11T19:31:37.379085Z","iopub.status.idle":"2023-07-11T19:31:37.855302Z","shell.execute_reply.started":"2023-07-11T19:31:37.379052Z","shell.execute_reply":"2023-07-11T19:31:37.854323Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\" ========== predicted ordered cell_ids ==========\")\nprint(out[\"prediction\"])","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:37.856871Z","iopub.execute_input":"2023-07-11T19:31:37.857494Z","iopub.status.idle":"2023-07-11T19:31:37.863461Z","shell.execute_reply.started":"2023-07-11T19:31:37.857442Z","shell.execute_reply":"2023-07-11T19:31:37.862464Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\" ========== predicted probabilities ==========\")\nprint(out[\"probs\"])","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:37.864871Z","iopub.execute_input":"2023-07-11T19:31:37.865456Z","iopub.status.idle":"2023-07-11T19:31:37.879961Z","shell.execute_reply.started":"2023-07-11T19:31:37.865419Z","shell.execute_reply":"2023-07-11T19:31:37.878870Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Since the pipeline appends begin-of-code-cell and end-of-code-cell to code cells, num_code_cell_pairs is one bigger than num_code_cells.\nNOTE: Here I use the strings \"Begin-of-CodeCell\" and \"End-of-CodeCell\" for ease of visualization, however, during the actual training and inference, an additional BERT special token \"[BLANK_CODE]\" is used to represent both of them. Don't use the strings \"Begin-of-CodeCell\" nor \"End-of-CodeCell\" for the input of the bert-tokenizer. ","metadata":{}},{"cell_type":"code","source":"# codes_with_sentinel = [\"[BLANK_CODE]\"] + list(pipeline_input[\"code\"][\"texts\"]) + [\"[BLANK_CODE]\"]\ncodes_with_sentinel = [\"Begin-of-CodeCell\"] + list(pipeline_input[\"code\"][\"texts\"]) + [\"End-of-CodeCell\"]\n\ncode_pairs = list()\nfor i in range(len(codes_with_sentinel)-1):\n    pair = [codes_with_sentinel[i], codes_with_sentinel[i+1]]\n    code_pairs.append(pair)\n\nprint(\"num_code_cells:\", len(pipeline_input[\"code\"][\"texts\"]))\nprint(\"num_code_cells_with_sentinel:\", len(codes_with_sentinel))\nprint(\"num_code_cell_pairs:\", len(code_pairs))\nprint()\nprint(\"num_markdown_cells:\", len(pipeline_input[\"markdown\"][\"texts\"]))\nprint()\nprint(\"shape of output probs(=[num_markdown_cells, num_code_cell_pairs]):\", out[\"probs\"].shape)","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:32:27.154721Z","iopub.execute_input":"2023-07-11T19:32:27.155083Z","iopub.status.idle":"2023-07-11T19:32:27.165863Z","shell.execute_reply.started":"2023-07-11T19:32:27.155053Z","shell.execute_reply":"2023-07-11T19:32:27.164468Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Check predicted positioning","metadata":{}},{"cell_type":"code","source":"TARGET_MARKDOWN_I = 0\n\nprint(\" ========== target markdown cell index ==========\")\nprint(TARGET_MARKDOWN_I)\nprint()\n\nprint(\" ========== target markdown cell text ==========\")\ntarget_markdown_cell = pipeline_input[\"markdown\"][\"texts\"][TARGET_MARKDOWN_I]\nprint(target_markdown_cell)\nprint()\n\nprint(\" ========== probs ==========\")\nprobs_for_target_markdown = out[\"probs\"][TARGET_MARKDOWN_I]\nfor i, prob in enumerate(probs_for_target_markdown):\n    print(f'{i}-th code-pair: {prob*100.0:5.02f}%')\nprint()\n\nprint(\" ========== insert to predicted position ==========\")\nselected_code_pair_i = probs_for_target_markdown.argmax().item()\npre_markdown_cell, post_markdown_cell = code_pairs[selected_code_pair_i]\n# selected_code_pair = [pre_markdown, post_markdown]\n\nprint(\" ---------- pre-markdown code_cell ----------\")\nprint(pre_markdown_cell)\nprint()\nprint(\" ---------- target markdown_cell ----------\")\nprint(target_markdown_cell)\nprint()\nprint(\" ---------- post-markdown code_cell ----------\")\nprint(post_markdown_cell)\n","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:37.895319Z","iopub.execute_input":"2023-07-11T19:31:37.896291Z","iopub.status.idle":"2023-07-11T19:31:37.910224Z","shell.execute_reply.started":"2023-07-11T19:31:37.896255Z","shell.execute_reply":"2023-07-11T19:31:37.909061Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## The output has representation vectors for code cell pairs and markdown cells.","metadata":{}},{"cell_type":"code","source":"print(\"code-pair_reprs(=[num_pairs, repr_dim]):\", out[\"code_reprs\"].shape)\nout[\"code_reprs\"]","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:37.913190Z","iopub.execute_input":"2023-07-11T19:31:37.914458Z","iopub.status.idle":"2023-07-11T19:31:37.927961Z","shell.execute_reply.started":"2023-07-11T19:31:37.914424Z","shell.execute_reply":"2023-07-11T19:31:37.926918Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"markdown_reprs(=[num_markdowns, repr_dim]):\", out[\"markdown_reprs\"].shape)\nout[\"markdown_reprs\"]","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:37.929676Z","iopub.execute_input":"2023-07-11T19:31:37.930390Z","iopub.status.idle":"2023-07-11T19:31:37.941394Z","shell.execute_reply.started":"2023-07-11T19:31:37.930357Z","shell.execute_reply":"2023-07-11T19:31:37.940514Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Head-less forwarding\nIf you only need representation vectors, use `pipe.get_network_output` with `reprs_only=True`.","metadata":{}},{"cell_type":"code","source":"from typing import List\ncode_texts:List[str] = notebook_data[\"code\"][\"texts\"]\nmarkdown_texts:List[str] = notebook_data[\"markdown\"][\"texts\"]\n\nwith torch.no_grad():\n    network_out = pipe.get_network_output(code_texts, markdown_texts, reprs_only=True)\n\nprint(\"code-pair_reps:\", network_out.code_reprs.shape)\nprint(\"markdown_reps:\", network_out.markdown_reprs.shape)","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:37.943004Z","iopub.execute_input":"2023-07-11T19:31:37.943673Z","iopub.status.idle":"2023-07-11T19:31:38.320012Z","shell.execute_reply.started":"2023-07-11T19:31:37.943639Z","shell.execute_reply":"2023-07-11T19:31:38.318856Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"------------","metadata":{}},{"cell_type":"markdown","source":"# Supplemental: step-by-step prediction\n\nWe can divide the pipeline process into tokenization and network-forwarding processes. This could be useful for the further fine-tuning.\n\nFor more infomation, see the source code of `ai4code_parallel_bert`.","metadata":{}},{"cell_type":"code","source":"bert_tokenizer = pipe.tokenizer\n\neach_code_max_length = pipe.each_code_max_length\nmarkdown_max_length = pipe.markdown_max_length\n\n# tokenize\ncode_texts = notebook_data[\"code\"][\"texts\"]\nmarkdown_texts = notebook_data[\"markdown\"][\"texts\"]\n\ncode_tokenized = bert_tokenizer(code_texts, add_special_tokens=False, truncation=True, max_length=each_code_max_length)[\"input_ids\"]\nmarkdown_tokenized = bert_tokenizer(markdown_texts, add_special_tokens=False, truncation=True, max_length=markdown_max_length)[\"input_ids\"]\n\n# forward\ncode_inputs, markdown_inputs = pipe.format_input(code_tokenized, markdown_tokenized) # this adds special tokens and begin/end-of-code-cell, then make pairs of code cells.\nwith torch.no_grad():\n    probs, cosines, code_reprs, markdown_reprs = pipe.network_forward(code_inputs, markdown_inputs, reprs_only=False)\n\n# make ordered cell_ids\npreds = probs.argmax(1).tolist()\ncode_cell_ids = notebook_data[\"code\"][\"cell_ids\"]\nmarkdown_cell_ids = notebook_data[\"markdown\"][\"cell_ids\"]\n\npred_order = [list() for _ in range(probs.shape[1])]\nfor code_i, code_cell_id in enumerate(code_cell_ids):\n    pred_order[code_i+1].append(code_cell_id)\nfor pred_position, markdown_cell_id in zip(preds, markdown_cell_ids):\n    pred_order[pred_position].append(markdown_cell_id)\npred_order = sum(pred_order, list())\n\noutput = {\n    \"prediction\": pred_order,\n    \"probs\": probs,\n    \"cosines\": cosines,\n    \"code_reprs\": code_reprs,\n    \"markdown_reprs\": markdown_reprs,\n}\n\nprint(\"probs:\", probs.shape)\nprint(\"code-pair_reprs:\", code_reprs.shape)\nprint(\"markdown_reprs:\", markdown_reprs.shape)","metadata":{"execution":{"iopub.status.busy":"2023-07-11T19:31:57.205774Z","iopub.execute_input":"2023-07-11T19:31:57.206178Z","iopub.status.idle":"2023-07-11T19:31:57.600754Z","shell.execute_reply.started":"2023-07-11T19:31:57.206146Z","shell.execute_reply":"2023-07-11T19:31:57.597975Z"},"trusted":true},"execution_count":null,"outputs":[]}]}