create_test_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e7219917",
   "metadata": {},
   "source": [
    "# Create test dataset, by excluding certain queries from training and augmented datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "747fe513",
   "metadata": {},
   "source": [
    "## Load training and augmented datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "e1e2c759",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import re\n",
    "df_train = pd.read_csv(\"train-data/sql_train.tsv\", sep='\\t')\n",
    "df_augment = pd.read_csv(\"train-data/data_augmentation_final.tsv\", sep='\\t')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6267d493",
   "metadata": {},
   "source": [
    "## Establish array of team names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "a59f9960",
   "metadata": {},
   "outputs": [],
   "source": [
    "team_array = [\n",
    "\"Atlanta Hawks\",\n",
    "\"Boston Celtics\",\n",
    "\"Cleveland Cavaliers\",\n",
    "\"New Orleans Pelicans\",\n",
    "\"Chicago Bulls\",\n",
    "\"Dallas Mavericks\",\n",
    "\"Denver Nuggets\",\n",
    "\"Golden State Warriors\",\n",
    "\"Houston Rockets\",\n",
    "\"Los Angeles Clippers\",\n",
    "\"Los Angeles Lakers\",\n",
    "\"Miami Heat\",\n",
    "\"Milwaukee Bucks\",\n",
    "\"Minnesota Timberwolves\",\n",
    "\"Brooklyn Nets\",\n",
    "\"New York Knicks\",\n",
    "\"Orlando Magic\",\n",
    "\"Indiana Pacers\",\n",
    "\"Philadelphia 76ers\",\n",
    "\"Phoenix Suns\",\n",
    "\"Portland Trail Blazers\",\n",
    "\"Sacramento Kings\",\n",
    "\"San Antonio Spurs\",\n",
    "\"Oklahoma City Thunder\",\n",
    "\"Toronto Raptors\",\n",
    "\"Utah Jazz\",\n",
    "\"Memphis Grizzlies\",\n",
    "\"Washington Wizards\",\n",
    "\"Detroit Pistons\",\n",
    "\"Charlotte Hornets\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87dbf2dd",
   "metadata": {},
   "source": [
    "## Define regex expression to build new columns in the dataframes with the team names masked out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "188544ef",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Which NBA teams were established after the year 2000? List their names and founding years, sorted from newest to oldest\n",
      "What is the most points the <TEAM> have scored at home?\n",
      "What is the second-highest number of points the <TEAM> have ever scored in a single home game?\n",
      "How many home games did the <TEAM> win in the 2017 season?\n",
      "What is the average number of assists by the <TEAM> in home wins?\n",
      "Which game had the highest total points scored by both teams when the <TEAM> played at home?\n",
      "What is the <TEAM>' largest lead in a home game during the 2016 season?\n",
      "How many times did the <TEAM> score over 120 points at home?\n",
      "What is the highest three-point percentage the <TEAM> achieved in an away game?\n",
      "What was the average points difference in home games won by the <TEAM>?\n"
     ]
    }
   ],
   "source": [
    "# Compile a regex pattern for matching any team name\n",
    "team_pattern = re.compile(\"|\".join(map(re.escape, sorted(team_array, key=len, reverse=True))))\n",
    "\n",
    "def normalize_query(query):\n",
    "    \"\"\"Replace any known team name with <TEAM>.\"\"\"\n",
    "    return team_pattern.sub(\"<TEAM>\", query)\n",
    "\n",
    "# Assuming `df_train` and `df_augmented` are your two dataframes\n",
    "df_train[\"template\"] = df_train[\"natural_query\"].apply(normalize_query)\n",
    "df_augment[\"template\"] = df_augment[\"natural_query\"].apply(normalize_query)\n",
    "\n",
    "# Test performance of regex masking\n",
    "count = 0\n",
    "for index, row in df_train.iterrows():\n",
    "    print(row['template'])\n",
    "    count += 1\n",
    "    if count == 10:\n",
    "        break\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9ff879f",
   "metadata": {},
   "source": [
    "## Create new dataframe to hold test data, delete entries from test set in train and augment dataframes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "ffb238d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1044\n",
      "786\n",
      "1870\n",
      "1367\n",
      "150\n"
     ]
    }
   ],
   "source": [
    "# Create new dataframes\n",
    "#df_test = pd.DataFrame(columns=['natural_query', 'sql_query', 'result'])\n",
    "#df_new_train = df_train.copy(deep=True)\n",
    "#df_new_augment = df_augment.copy(deep=True)\n",
    "\n",
    "# Sample some test queries from the original training set\n",
    "test_samples = df_train.sample(n=150)  # Or whatever number you want\n",
    "\n",
    "# Get the set of template forms of the test samples\n",
    "test_templates = set(test_samples[\"template\"])\n",
    "\n",
    "# Filter out those entries from the train and augmented sets\n",
    "df_train_filtered = df_train[~df_train[\"template\"].isin(test_templates)].drop(columns=[\"template\"])\n",
    "df_augment_filtered = df_augment[~df_augment[\"template\"].isin(test_templates)].drop(columns=[\"template\"])\n",
    "\n",
    "print(len(df_train))\n",
    "print(len(df_train_filtered))\n",
    "print(len(df_augment))\n",
    "print(len(df_augment_filtered))\n",
    "print(len(test_samples))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84896991",
   "metadata": {},
   "source": [
    "## Check distribution of new datasets to ensure they match original"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "9d3587ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "36\n"
     ]
    }
   ],
   "source": [
    "count = 0\n",
    "for index, row in test_samples.iterrows():\n",
    "    if len(row['sql_query']) <= 90:\n",
    "        count += 1\n",
    "\n",
    "print(count)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2484fabf",
   "metadata": {},
   "source": [
    "## Save new train and test dataframes to tsv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c546ba58",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "150\n",
      "2213\n"
     ]
    }
   ],
   "source": [
    "df = pd.concat([df_train_filtered, df_augment_filtered])\n",
    "print(len(test_samples))\n",
    "print(len(df))\n",
    "test_samples.to_csv('./train-data/test_set.tsv', index=False, sep='\\t', columns=['natural_query', 'sql_query', 'result'])\n",
    "df.to_csv('./train-data/train_set.tsv', index=False, sep='\\t', columns=['natural_query', 'sql_query', 'result'])"
   ]
  }
 ],
 "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}