{"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":"!pip install sentence-transformers -qq","metadata":{"execution":{"iopub.status.busy":"2022-05-12T12:01:38.08372Z","iopub.execute_input":"2022-05-12T12:01:38.084046Z","iopub.status.idle":"2022-05-12T12:01:52.725097Z","shell.execute_reply.started":"2022-05-12T12:01:38.08401Z","shell.execute_reply":"2022-05-12T12:01:52.723882Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This is a sample demo on how to use [Flax Code Embeddings](https://huggingface.co/flax-sentence-embeddings/st-codesearch-distilroberta-base) with [Google AI4Code](https://www.kaggle.com/competitions/AI4Code/) dataset. Like with any other SBert model, this is also finetunable to suite the purpose of this competition's task. The rest of the code is a sample usage demo to evaluate this model out of the box.\n\nFrom a first impression, it seems to do above average with English markdown cells, but not so much with multilingual corpus. This is expected since it is finetuned over Roberta, which is trained specifically on English corpora.","metadata":{}},{"cell_type":"code","source":"from sentence_transformers import SentenceTransformer, util\nimport pandas as pd, numpy as np, json\nfrom colorama import Fore, Back, Style\n\nmodel = SentenceTransformer(\"flax-sentence-embeddings/st-codesearch-distilroberta-base\")\nTRAIN_BASE_DIR = \"../input/AI4Code/train\"\n\ndf = pd.read_csv(\"../input/AI4Code/train_orders.csv\")","metadata":{"execution":{"iopub.status.busy":"2022-05-12T12:01:52.727248Z","iopub.execute_input":"2022-05-12T12:01:52.727546Z","iopub.status.idle":"2022-05-12T12:02:30.667477Z","shell.execute_reply.started":"2022-05-12T12:01:52.72751Z","shell.execute_reply":"2022-05-12T12:02:30.66644Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"row = df.sample().iloc[0]\ncid = row[\"id\"]\ncells = row[\"cell_order\"].split()\n\nwith open(f\"{TRAIN_BASE_DIR}/{cid}.json\") as f:\n    dat = json.load(f)\n\ncodes =  [dat[\"source\"][cell] for cell in cells if dat[\"cell_type\"][cell]==\"code\"]\nqueries = [dat[\"source\"][cell] for cell in cells if dat[\"cell_type\"][cell]!=\"code\"]\ncode_ids =  [cell for cell in cells if dat[\"cell_type\"][cell]==\"code\"]\nquery_ids = [cell for cell in cells if dat[\"cell_type\"][cell]!=\"code\"]","metadata":{"execution":{"iopub.status.busy":"2022-05-12T12:02:30.668924Z","iopub.execute_input":"2022-05-12T12:02:30.669171Z","iopub.status.idle":"2022-05-12T12:02:30.698122Z","shell.execute_reply.started":"2022-05-12T12:02:30.669139Z","shell.execute_reply":"2022-05-12T12:02:30.697491Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Ordered cells display","metadata":{}},{"cell_type":"code","source":"for cell in cells: #Ordered\n    content = dat[\"source\"][cell]\n    colour = Fore.GREEN if dat[\"cell_type\"][cell]==\"code\" else Fore.BLUE\n    print(f\"Cell id: {cell}\")\n    print(colour + content + Style.RESET_ALL)\n    print(\"$\"*50)","metadata":{"execution":{"iopub.status.busy":"2022-05-12T12:02:30.699929Z","iopub.execute_input":"2022-05-12T12:02:30.700358Z","iopub.status.idle":"2022-05-12T12:02:30.740618Z","shell.execute_reply.started":"2022-05-12T12:02:30.700322Z","shell.execute_reply":"2022-05-12T12:02:30.725676Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Predict closest code cell for a list of markdown cells","metadata":{}},{"cell_type":"code","source":"code_emb = model.encode(codes, convert_to_tensor=True)\nquery_emb = model.encode(queries, convert_to_tensor=True)\nquery_hits = util.semantic_search(query_emb, code_emb)\n\nop_dat = [{\"score\": hits[0][\"score\"], \n           \"code_match\": code_ids[hits[0][\"corpus_id\"]], \n           \"markdown\": qid} for qid, hits in zip(query_ids, query_hits)]\nfor i, dct in enumerate(op_dat):\n    dct[\"abs_dist\"] = abs(cells.index(dct[\"code_match\"]) - cells.index(dct[\"markdown\"]))\n\nprint(json.dumps(op_dat, indent=2))","metadata":{"execution":{"iopub.status.busy":"2022-05-12T12:02:30.742099Z","iopub.execute_input":"2022-05-12T12:02:30.742441Z","iopub.status.idle":"2022-05-12T12:02:36.223104Z","shell.execute_reply.started":"2022-05-12T12:02:30.742383Z","shell.execute_reply":"2022-05-12T12:02:36.222496Z"},"trusted":true},"execution_count":null,"outputs":[]}]}