0. Generating Splits.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4e5bbaa",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-12-16T13:22:40.102324Z",
     "iopub.status.busy": "2022-12-16T13:22:40.101661Z",
     "iopub.status.idle": "2022-12-16T13:22:41.670741Z",
     "shell.execute_reply": "2022-12-16T13:22:41.669120Z"
    },
    "papermill": {
     "duration": 1.58498,
     "end_time": "2022-12-16T13:22:41.674272",
     "exception": false,
     "start_time": "2022-12-16T13:22:40.089292",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import matplotlib\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm.notebook import tqdm\n",
    "from sklearn.model_selection import StratifiedGroupKFold, GroupKFold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ac9a9a3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-12-16T13:22:41.700052Z",
     "iopub.status.busy": "2022-12-16T13:22:41.699543Z",
     "iopub.status.idle": "2022-12-16T13:22:41.705438Z",
     "shell.execute_reply": "2022-12-16T13:22:41.703955Z"
    },
    "papermill": {
     "duration": 0.023327,
     "end_time": "2022-12-16T13:22:41.708354",
     "exception": false,
     "start_time": "2022-12-16T13:22:41.685027",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "DATA_PATH = \"../raw_data/\" # raw_data路径"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c38c14b5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-12-16T13:22:41.731837Z",
     "iopub.status.busy": "2022-12-16T13:22:41.731417Z",
     "iopub.status.idle": "2022-12-16T13:23:05.855527Z",
     "shell.execute_reply": "2022-12-16T13:23:05.854203Z"
    },
    "papermill": {
     "duration": 24.139111,
     "end_time": "2022-12-16T13:23:05.858134",
     "exception": false,
     "start_time": "2022-12-16T13:22:41.719023",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "topics = pd.read_csv(DATA_PATH + \"topics.csv\")\n",
    "content = pd.read_csv(DATA_PATH + \"content.csv\")\n",
    "correlations = pd.read_csv(DATA_PATH + \"correlations.csv\")\n",
    "\n",
    "topics = topics[topics.has_content==True] # 只保留有content的topic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb493013",
   "metadata": {},
   "source": [
    "## Split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9c26411",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 将所有category==\"source\"的topic，都划分到train fold中，因为隐藏测试集中没有category==\"source\"的topic\n",
    "topics_train = topics[topics['category'] == \"source\"][[\"id\"]] \n",
    "topics_train[\"fold\"] = \"train\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26e43ad6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-12-16T13:23:12.222085Z",
     "iopub.status.busy": "2022-12-16T13:23:12.221455Z",
     "iopub.status.idle": "2022-12-16T13:23:12.465709Z",
     "shell.execute_reply": "2022-12-16T13:23:12.464474Z"
    },
    "papermill": {
     "duration": 0.264501,
     "end_time": "2022-12-16T13:23:12.468560",
     "exception": false,
     "start_time": "2022-12-16T13:23:12.204059",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# valid df\n",
    "topics_val = topics[topics['category'] != \"source\"].reset_index(drop=True)\n",
    "\n",
    "# StratifiedGroupKFold\n",
    "sgkf = StratifiedGroupKFold(random_state=1773,\n",
    "                            n_splits=4,\n",
    "                            shuffle=True)\n",
    "\n",
    "# 根据language进行分层抽样，根据id进行分组\n",
    "split_idxs = list(sgkf.split(topics_val[\"id\"], topics_val[\"language\"], groups=topics_val[\"id\"]))[0]\n",
    "\n",
    "split_idxs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d10b66af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建训练集\n",
    "topics_add_train = topics_val.iloc[split_idxs[0]].reset_index(drop=True)[[\"id\"]]\n",
    "topics_add_train[\"fold\"] = \"train\"\n",
    "\n",
    "topics_train = pd.concat([topics_train, topics_add_train], ignore_index=True)\n",
    "\n",
    "# 创建验证集\n",
    "topics_holdout = topics_val.iloc[split_idxs[1]].reset_index(drop=True)[[\"id\"]]\n",
    "topics_holdout[\"fold\"] = \"test\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a94b0d77",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 查看训练集的language分布\n",
    "topics[topics.id.isin(topics_train.id)].language.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c559862e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# \n",
    "topics[topics.id.isin(topics_holdout.id)].language.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b4aadf3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-12-16T13:23:15.013233Z",
     "iopub.status.busy": "2022-12-16T13:23:15.012768Z",
     "iopub.status.idle": "2022-12-16T13:23:15.020818Z",
     "shell.execute_reply": "2022-12-16T13:23:15.019953Z"
    },
    "papermill": {
     "duration": 0.03394,
     "end_time": "2022-12-16T13:23:15.023192",
     "exception": false,
     "start_time": "2022-12-16T13:23:14.989252",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 保存完整数据集\n",
    "split_df = pd.concat([topics_train, topics_holdout], ignore_index=True)\n",
    "split_df.to_csv('train_test_splits.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47cb14bf",
   "metadata": {
    "papermill": {
     "duration": 0.021488,
     "end_time": "2022-12-16T13:23:15.131659",
     "exception": false,
     "start_time": "2022-12-16T13:23:15.110171",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Done !"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 45.346996,
   "end_time": "2022-12-16T13:23:16.176659",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2022-12-16T13:22:30.829663",
   "version": "2.3.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}