Multiclass_Logistic_Regression_Gender/train_dev_test_split.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "66be5e0f",
   "metadata": {},
   "source": [
    "### Overview\n",
    "\n",
    "In this script, I first split the CCNC (Comprehensive Chinese Name Corpus) into train/dev/test sets based on the ratio of 8:1:1, which translates into 2926487: 365811: 365811 in terms of sample size.\n",
    "\n",
    "As only the first names will be used to train the logistic regression model, both the train set and dev set are then deduplicated based on first names and gender (so that the model can be more efficiently trained). This process reduces the size of train set from 2926487 to 633857, and the size of dev set from 365811 to 167528. The finally saved train/dev sets only contain first names and related genders.\n",
    "\n",
    "However, in the test set, the full names and related genders are instead saved and it is thus deduplicated based on full names and gender (if I do it for first name + gender, the sample will be reduced from 365811 to 167849). The full names are saved because there will be a rule-based filter to retrieve the first names from given names and I want to test whether that works well. \n",
    "\n",
    "Therefore, in actuality, the ratio among the train/dev/test sets is 633857: 167528: 167849, which roughly equals to 6: 2: 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2feaa823",
   "metadata": {},
   "outputs": [],
   "source": [
    "from random import shuffle, seed\n",
    "\n",
    "\n",
    "def readFile(filepath='data/ccnc.txt'):\n",
    "    '''Read a file given its file path.\n",
    "    \n",
    "    Paras:\n",
    "        file: filepath. Defaults to 'ccnc.txt'. Please\n",
    "            ensure this file is within the current directory. \n",
    "    \n",
    "    '''\n",
    "    data = []\n",
    "    f = open(filepath, 'r')\n",
    "    header = next(f)\n",
    "    for line in f:\n",
    "        data.append(line)\n",
    "    return data\n",
    "\n",
    "\n",
    "def train_dev_test_split(data, train=0.6, dev=0.2, test=0.2, seed_idx=5):\n",
    "    '''\n",
    "    Split ccnc.txt into train, dev and test sets with a predefined ratio.\n",
    "    \n",
    "    Paras:\n",
    "        train, dev, test: respective ratio for the train, dev and test sets. \n",
    "            Default to 0.6, 0.2, 0.2 respectively. \n",
    "        seed_idx: Int. Defaults to 5 (a random picked seed). \n",
    "    '''\n",
    "    \n",
    "    seed(seed_idx)\n",
    "    shuffle(data)\n",
    "    length = len(data)\n",
    "    boundary1 = round(length * train)\n",
    "    boundary2 = round(length * (train + dev))\n",
    "    \n",
    "    # return the tran_ds, dev_ds, test_ds \n",
    "    return data[:boundary1], data[boundary1: boundary2], data[boundary2:]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c9a54bc3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " train set size: 2926487\n",
      " dev set size: 365811\n",
      " test set size: 365811\n"
     ]
    }
   ],
   "source": [
    "data = readFile()\n",
    "train_ds, dev_ds, test_ds = train_dev_test_split(data, 0.8, 0.1, 0.1)\n",
    "# debug to see whether the spliting is working\n",
    "assert(len(data) == len(train_ds) + len(dev_ds) + len(test_ds))\n",
    "print(f' train set size: {len(train_ds)}\\n dev set size: {len(dev_ds)}\\n test set size: {len(test_ds)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b3ad0c8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# deduplicate the train set and dev set based on first name and gender\n",
    "\n",
    "from collections import Counter\n",
    "\n",
    "\n",
    "def deduplicate(data, is_test=False):\n",
    "    '''This function deduplicates the train set and dev set based on first name and gender\n",
    "    and returns dict_lists that only contain first names and genders. If the data is the test set,\n",
    "    this function will return a dict_dict that contains only full names and genders. \n",
    "    Paras:\n",
    "        data: list --> train_ds or dev_ds\n",
    "    '''\n",
    "    count = Counter()\n",
    "    if not is_test:\n",
    "        for example in data:\n",
    "            example = example.split('\\t')\n",
    "            fname, gender = example[1], example[3] # fname = first name\n",
    "            # deduplicate based on first name and gender\n",
    "            count[fname + '\\t' + gender] = 1\n",
    "    else:\n",
    "        for example in data:\n",
    "            example = example.split('\\t')\n",
    "            Fname, gender = example[2], example[3] # Fname = full name\n",
    "            # deduplicate based on full name and gender\n",
    "            count[Fname + '\\t' + gender] = 1\n",
    "    return count.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e1ffdb25",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " new train set: 633857\n",
      " new dev set: 167528\n",
      " new test set: 365810\n"
     ]
    }
   ],
   "source": [
    "train_ds_new = deduplicate(train_ds)\n",
    "dev_ds_new = deduplicate(dev_ds)\n",
    "test_ds_new = deduplicate(test_ds, is_test=True)\n",
    "print(f' new train set: {len(train_ds_new)}\\n new dev set: {len(dev_ds_new)}\\n new test set: {len(test_ds_new)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "63765069",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fileWriter(data, file_name):\n",
    "    '''Write a list of name examples back into a txt file and save in \n",
    "    the current directory if the full path is not given in the file_name.\n",
    "    \n",
    "    Paras:\n",
    "        data: list\n",
    "            a list of name examples; each example contains \n",
    "            last name, first name, full name and gender\n",
    "        file_name: str\n",
    "    '''\n",
    "    file_name = file_name if file_name.endswith('.txt') else file_name + '.txt'\n",
    "    tmp = '{}\\t{}'\n",
    "    with open(file_name, 'w') as f:\n",
    "        f.write(tmp.format('name', 'gender\\n'))\n",
    "        f.write(''.join(data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d4bde832",
   "metadata": {},
   "outputs": [],
   "source": [
    "fileWriter(train_ds_new, 'data/train_ds.txt')\n",
    "fileWriter(dev_ds_new, 'data/dev_ds.txt')\n",
    "fileWriter(test_ds_new, 'data/test_ds.txt')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}