{"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":"!cp -r ../input/ai4code-source/* ./\n\nimport tensorflow as tf\nimport pandas as pd\nimport random\nimport numpy as np\n\nfrom model.model import Model\nfrom data import Dataset\n\nTOKENIZERS_PARALLELISM = False","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:36:44.033452Z","iopub.execute_input":"2022-10-01T15:36:44.033911Z","iopub.status.idle":"2022-10-01T15:36:51.799161Z","shell.execute_reply.started":"2022-10-01T15:36:44.033831Z","shell.execute_reply":"2022-10-01T15:36:51.798113Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Constant","metadata":{}},{"cell_type":"code","source":"MAX_LEN = 64\nMAX_CELL = 128\nMODEL_PATH = 'microsoft/codebert-base'\nD_MODEL = 768\nNUM_TEST = 7000\nDATA_DIR = \"../input/AI4Code\"\nCELL_PAD = 0\nBATCH_SIZE = 2","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:36:51.801270Z","iopub.execute_input":"2022-10-01T15:36:51.801965Z","iopub.status.idle":"2022-10-01T15:36:51.815051Z","shell.execute_reply.started":"2022-10-01T15:36:51.801927Z","shell.execute_reply":"2022-10-01T15:36:51.813093Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data preparation","metadata":{}},{"cell_type":"markdown","source":"Import data","metadata":{}},{"cell_type":"code","source":"test_df = pd.read_csv(\"../input/ai4code-split-dataset/data/test.csv\")","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:36:51.816975Z","iopub.execute_input":"2022-10-01T15:36:51.817705Z","iopub.status.idle":"2022-10-01T15:36:54.177813Z","shell.execute_reply.started":"2022-10-01T15:36:51.817669Z","shell.execute_reply":"2022-10-01T15:36:54.176844Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Get a subset of notebooks for training","metadata":{}},{"cell_type":"code","source":"nb_ids = test_df.groupby('id').groups.keys()\n\nn = min(test_df.id.nunique(), NUM_TEST)\nrandom.seed(42)\nrandom_ids = random.sample(nb_ids, n)\n\ntest_df = test_df[test_df.id.isin(random_ids)].reset_index(drop=True)","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:36:54.180293Z","iopub.execute_input":"2022-10-01T15:36:54.181114Z","iopub.status.idle":"2022-10-01T15:36:54.461343Z","shell.execute_reply.started":"2022-10-01T15:36:54.181073Z","shell.execute_reply":"2022-10-01T15:36:54.460363Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Sort dataset by ```pct_rank```","metadata":{}},{"cell_type":"code","source":"test_df = test_df.sort_values(['id', 'pct_rank'])","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:36:54.462810Z","iopub.execute_input":"2022-10-01T15:36:54.463716Z","iopub.status.idle":"2022-10-01T15:36:54.474962Z","shell.execute_reply.started":"2022-10-01T15:36:54.463676Z","shell.execute_reply":"2022-10-01T15:36:54.473991Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Create tensor dataset for training","metadata":{}},{"cell_type":"code","source":"d = Dataset(\n    model_path=MODEL_PATH, \n    max_len=MAX_LEN, \n    num_cells=MAX_CELL,  \n    batch_size=BATCH_SIZE, \n    cell_pad=CELL_PAD\n)\n\nprocessed_test_df, test = d.build_dataset(test_df)","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:36:54.476696Z","iopub.execute_input":"2022-10-01T15:36:54.477050Z","iopub.status.idle":"2022-10-01T15:37:09.826245Z","shell.execute_reply.started":"2022-10-01T15:36:54.477017Z","shell.execute_reply":"2022-10-01T15:37:09.825332Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"# Input definitions\ninput_ids = tf.keras.layers.Input(shape=(MAX_CELL, MAX_LEN), name='input_ids', dtype='int32')\nattention_mask = tf.keras.layers.Input(shape=(MAX_CELL, MAX_LEN), name='attention_mask', dtype='int32')\ncell_features = tf.keras.layers.Input(shape=(MAX_CELL, 2), name='cell_features', dtype='float32')\ncell_mask = tf.keras.layers.Input(shape=(MAX_CELL), name='cell_mask', dtype='float32')\n\n# Model initialization\nstacked_transformer = Model(\n    model_path=MODEL_PATH, \n    d_model=D_MODEL, \n    n_heads=6, \n    num_cells=MAX_CELL,\n    dropout_pos=0.,\n    dropout_trans=0., \n    eps=1e-6, \n    d_ff_trans=64, \n    ff_activation=\"relu\", \n    n_layers=4,\n    d_ff_pointwise=64, \n    dropout_pointwise=0.\n)\n\n# Define the model\nmodel = tf.keras.Model(\n    inputs=[input_ids, attention_mask, cell_features, cell_mask], \n    outputs=stacked_transformer.call(input_ids, attention_mask, cell_features, cell_mask, False)\n)","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:38:57.866088Z","iopub.execute_input":"2022-10-01T15:38:57.866484Z","iopub.status.idle":"2022-10-01T15:39:02.378377Z","shell.execute_reply.started":"2022-10-01T15:38:57.866450Z","shell.execute_reply":"2022-10-01T15:39:02.377425Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.load_weights(\"../input/ai4codemodel/saved_model/\")","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:39:02.380131Z","iopub.execute_input":"2022-10-01T15:39:02.380512Z","iopub.status.idle":"2022-10-01T15:39:02.954371Z","shell.execute_reply.started":"2022-10-01T15:39:02.380458Z","shell.execute_reply":"2022-10-01T15:39:02.953283Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Model summary\nmodel.summary()","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:39:02.955840Z","iopub.execute_input":"2022-10-01T15:39:02.956308Z","iopub.status.idle":"2022-10-01T15:39:02.978792Z","shell.execute_reply.started":"2022-10-01T15:39:02.956270Z","shell.execute_reply":"2022-10-01T15:39:02.977919Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Metric","metadata":{}},{"cell_type":"code","source":"def count_inversions_slowly(ranks):\n    inversions = 0\n    size = len(ranks)\n    for i in range(size):\n        for j in range(i+1, size):\n            if ranks[i] > ranks[j]:\n                total += 1\n    return total\n\nfrom bisect import bisect\n\n\n# Actually O(N^2), but fast in practice for our data\ndef count_inversions(a):\n    inversions = 0\n    sorted_so_far = []\n    for i, u in enumerate(a):  # O(N)\n        j = bisect(sorted_so_far, u)  # O(log N)\n        inversions += i - j\n        sorted_so_far.insert(j, u)  # O(N)\n    return inversions\n\ndef kendall_tau(ground_truth, predictions):\n    total_inversions = 0  # total inversions in predicted ranks across all instances\n    total_2max = 0  # maximum possible inversions across all instances\n    for gt, pred in zip(ground_truth, predictions):\n        ranks = [gt.index(x) for x in pred]  # rank predicted order in terms of ground truth\n        total_inversions += count_inversions(ranks)\n        n = len(gt)\n        total_2max += n * (n - 1)\n    return 1 - 4 * total_inversions / total_2max","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:39:02.980599Z","iopub.execute_input":"2022-10-01T15:39:02.980861Z","iopub.status.idle":"2022-10-01T15:39:02.991180Z","shell.execute_reply.started":"2022-10-01T15:39:02.980837Z","shell.execute_reply":"2022-10-01T15:39:02.990200Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Inference","metadata":{}},{"cell_type":"markdown","source":"Reconstruction function","metadata":{}},{"cell_type":"code","source":"def map_pred_pct_rank(df, pred_pct_rank):\n    \n    def hehe(group, pred_pct_rank, ids):\n        id_nb = group.id.iloc[0]\n        id_count = np.where(ids == id_nb)[0][0]\n        \n        pct_rank = pred_pct_rank[id_count]\n        \n        group_len = len(group)\n        group[\"pred_pct_rank\"] = pct_rank[:group_len]\n        return group\n\n    temp = df.copy()\n    temp[\"pred_pct_rank\"] = 0.\n    \n    ids = temp.id.unique()\n    \n    groups = temp.groupby(\"id\").apply(hehe, pred_pct_rank, ids).reset_index(drop=True)\n    \n    groups[\"pred_pct_rank\"] = np.where(groups[\"cell_type\"] == \"code\", groups[\"pct_rank\"], groups[\"pred_pct_rank\"])\n    \n    groups = groups.sort_values(\"pred_pct_rank\").groupby(\"id\")[\"cell_id\"].apply(lambda x: \" \".join(x)).reset_index()\n    \n    return groups","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:39:02.992703Z","iopub.execute_input":"2022-10-01T15:39:02.993395Z","iopub.status.idle":"2022-10-01T15:39:03.003985Z","shell.execute_reply.started":"2022-10-01T15:39:02.993355Z","shell.execute_reply":"2022-10-01T15:39:03.003144Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Infer","metadata":{}},{"cell_type":"code","source":"preds = model.predict(test, verbose=1)","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:39:03.005369Z","iopub.execute_input":"2022-10-01T15:39:03.006456Z","iopub.status.idle":"2022-10-01T15:39:47.611878Z","shell.execute_reply.started":"2022-10-01T15:39:03.006416Z","shell.execute_reply":"2022-10-01T15:39:47.610854Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"temp = processed_test_df.merge(test_df, on=[\"id\", \"cell_id\"])\npred_rank = map_pred_pct_rank(temp, preds)\npred_rank.head()","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:39:47.613432Z","iopub.execute_input":"2022-10-01T15:39:47.614030Z","iopub.status.idle":"2022-10-01T15:39:47.720204Z","shell.execute_reply.started":"2022-10-01T15:39:47.613991Z","shell.execute_reply":"2022-10-01T15:39:47.719216Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# predictions=a.sort_values(\"pred_pct_rank\").groupby(\"id\")[\"cell_id\"].apply(lambda x: \" \".join(x)).reset_index()\nground_truth = temp.sort_values(\"pct_rank\").groupby(\"id\")[\"cell_id\"].apply(lambda x: \" \".join(x)).reset_index()\nground_truth.head()","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:39:47.721339Z","iopub.execute_input":"2022-10-01T15:39:47.721681Z","iopub.status.idle":"2022-10-01T15:39:47.740790Z","shell.execute_reply.started":"2022-10-01T15:39:47.721653Z","shell.execute_reply":"2022-10-01T15:39:47.739961Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gt = list(map(lambda x: x.split(\" \"), ground_truth.cell_id.values))\npre = list(map(lambda x: x.split(\" \"), pred_rank.cell_id.values))","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:39:47.742204Z","iopub.execute_input":"2022-10-01T15:39:47.742588Z","iopub.status.idle":"2022-10-01T15:39:47.748453Z","shell.execute_reply.started":"2022-10-01T15:39:47.742554Z","shell.execute_reply":"2022-10-01T15:39:47.747376Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Model Kendal Tau score","metadata":{}},{"cell_type":"code","source":"kendall_tau(gt, pre)","metadata":{"execution":{"iopub.status.busy":"2022-10-01T15:39:47.751616Z","iopub.execute_input":"2022-10-01T15:39:47.752072Z","iopub.status.idle":"2022-10-01T15:39:47.765909Z","shell.execute_reply.started":"2022-10-01T15:39:47.752039Z","shell.execute_reply":"2022-10-01T15:39:47.764828Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}