{"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":"import random\nimport multiprocess as mp\nfrom tqdm import tqdm\nimport pandas as pd\nimport numpy as np\nimport json\nfrom sklearn.model_selection import GroupShuffleSplit\nimport os\nfrom pathlib import Path\n\n\nSEED = 42\nPICKLE_PROTOCOL = 4","metadata":{"execution":{"iopub.status.busy":"2022-12-18T04:45:38.971686Z","iopub.execute_input":"2022-12-18T04:45:38.972260Z","iopub.status.idle":"2022-12-18T04:45:39.586521Z","shell.execute_reply.started":"2022-12-18T04:45:38.972151Z","shell.execute_reply":"2022-12-18T04:45:39.585298Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"non_en_ids = pd.read_pickle(\"/kaggle/input/ai4code-non-english-notebooks/non_en_ids.pkl\")\nnon_en_ids","metadata":{"execution":{"iopub.status.busy":"2022-12-18T04:45:39.588368Z","iopub.execute_input":"2022-12-18T04:45:39.588962Z","iopub.status.idle":"2022-12-18T04:45:39.624304Z","shell.execute_reply.started":"2022-12-18T04:45:39.588920Z","shell.execute_reply":"2022-12-18T04:45:39.623249Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def read_nb(path):\n    return (\n        pd.read_json(\n            path,\n            dtype={'cell_type': 'category', \n                'source': 'str'}\n        )\n        .assign(id=path.stem)\n        .rename_axis('cell_id')\n    )\n\n\ndef process_order(raw_data_dir):\n    df_orders = pd.read_csv(raw_data_dir / 'train_orders.csv')\n    df_orders['cell_order'] = df_orders['cell_order'].str.split()\n    df_orders = df_orders.explode('cell_order')\n    df_orders['rank'] = df_orders.groupby('id')['cell_order'].cumcount()\n    df_orders['pct_rank'] = (\n        df_orders['rank'] / df_orders.groupby('id')['cell_order'].transform('count')\n    )\n    df_orders.rename(columns={'cell_order': 'cell_id'}, inplace=True)\n    return df_orders\n\n\ndef obtain_nb_info(df_merge, raw_data_dir):\n    df_merge['n_code_cells'] = (df_merge['cell_type'] == 'code').astype(np.int8)\n    df_merge['n_md_cells'] = (df_merge['cell_type'] == 'markdown').astype(np.int8)\n\n    df_nb = df_merge.groupby('id', as_index=False).agg({\n        'cell_id': 'count',\n        'n_code_cells': 'sum',\n        'n_md_cells': 'sum'\n    }).rename(columns={'cell_id': 'n_cells'})\n    df_nb['md_pct'] = df_nb['n_md_cells'] / df_nb['n_cells']\n    df_ancestors = pd.read_csv(raw_data_dir / 'train_ancestors.csv', index_col='id')\n    df_nb['ancestor_id'] = df_nb['id'].map(df_ancestors['ancestor_id'])\n\n    # A dict for all notebook metadata\n    nb_meta_data = df_nb.drop('ancestor_id', axis=1).set_index('id').to_dict(orient='index')\n    for d in nb_meta_data.values():\n        d['n_code_cells'] = int(d['n_code_cells'])\n        d['n_md_cells'] = int(d['n_md_cells'])\n    return df_nb, nb_meta_data\n\n\ndef train_val_split(nb_meta_data, df_nb, val_size):\n    splitter = GroupShuffleSplit(n_splits=1, test_size=val_size, random_state=SEED)\n    train_ind, val_ind = next(splitter.split(df_nb, groups=df_nb[\"ancestor_id\"]))\n    train_df = df_nb.loc[train_ind, 'id'].reset_index(drop=True)\n    val_df = df_nb.loc[val_ind, 'id'].reset_index(drop=True)\n\n    return train_df, val_df\n\n\ndef preprocess(non_en_ids):\n    raw_data_dir = Path(\"/kaggle/input/AI4Code\")\n    train_paths = list((raw_data_dir / 'train').glob('*.json'))\n    train_paths = np.array(train_paths)\n    \n    random.seed(SEED)\n    random.shuffle(train_paths)\n\n\n    with mp.Pool(mp.cpu_count()) as p:\n        notebooks_train = list(\n            p.map(\n                read_nb, \n                tqdm(train_paths, desc='Reading train notebooks')\n            )\n        )\n\n    df = pd.concat(notebooks_train).reset_index()\n\n    idx_in = np.isin(df[\"id\"].values, non_en_ids.values)\n    df = df.iloc[~idx_in].reset_index(drop=True)\n\n    df['is_code'] = (df['cell_type'] == 'code').astype(np.int8)\n    df['pos'] = df.groupby('id')['cell_id'].cumcount() + 1  # [1:TOTAL_MAX_CELLS]\n    # dummy start has 0 rel_pos, code cells have pos/n_code_cells => # last code cells, rel_pos = 1\n    df['rel_pos'] = df['pos'] / df.groupby('id')['is_code'].transform('sum')\n    df.loc[df['cell_type'] == 'markdown', 'rel_pos'] = 0.\n\n    df_orders = process_order(raw_data_dir)\n    df_merge = df.merge(df_orders, how='left', on=['id', 'cell_id'])\n    df_merge = df_merge[[\n        'id', \n        'cell_id', \n        'cell_type', \n        'is_code',\n        'pos', \n        'rel_pos',\n        'source', \n        'rank', \n        'pct_rank', \n    ]]\n\n    df_nb, nb_meta_data = obtain_nb_info(df_merge, raw_data_dir)\n\n    _, val_ids = train_val_split(nb_meta_data, df_nb, 0.1)\n\n    return val_ids\n\n\n\nval_ids = preprocess(non_en_ids)","metadata":{"execution":{"iopub.status.busy":"2022-12-18T04:53:16.591991Z","iopub.execute_input":"2022-12-18T04:53:16.592471Z","iopub.status.idle":"2022-12-18T05:04:05.137717Z","shell.execute_reply.started":"2022-12-18T04:53:16.592433Z","shell.execute_reply":"2022-12-18T05:04:05.135975Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(len(val_ids.values))","metadata":{"execution":{"iopub.status.busy":"2022-12-18T05:04:44.519839Z","iopub.execute_input":"2022-12-18T05:04:44.520281Z","iopub.status.idle":"2022-12-18T05:04:44.527248Z","shell.execute_reply.started":"2022-12-18T05:04:44.520243Z","shell.execute_reply":"2022-12-18T05:04:44.525895Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"raw_data_dir = Path(\"/kaggle/input/AI4Code\")\ndata_paths = list((raw_data_dir / 'train').glob('*.json'))\ndata_paths = np.array(data_paths)","metadata":{"execution":{"iopub.status.busy":"2022-12-18T05:04:57.349599Z","iopub.execute_input":"2022-12-18T05:04:57.350072Z","iopub.status.idle":"2022-12-18T05:04:58.071707Z","shell.execute_reply.started":"2022-12-18T05:04:57.350032Z","shell.execute_reply":"2022-12-18T05:04:58.070411Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"vfunc = np.vectorize(lambda x: str(x).split(\"/\")[-1].split(\".\")[0])\nids = vfunc(data_paths)\nidx_in = np.isin(ids, val_ids.values)\n\ncurr_paths = data_paths[idx_in]","metadata":{"execution":{"iopub.status.busy":"2022-12-18T05:04:58.074704Z","iopub.execute_input":"2022-12-18T05:04:58.075225Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cd /kaggle/working/\n!rm -rf test\n!mkdir test","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"vfunc = np.vectorize(lambda x: \"/kaggle/working/test/{id}.json\".format(id=x))\nnew_paths = vfunc(val_ids.values)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\n\nfor i in range(len(curr_paths)):\n    shutil.copyfile(curr_paths[i], new_paths[i])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"val_ids.to_pickle(\"/kaggle/working/test_ids.pkl\", protocol=PICKLE_PROTOCOL)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}