updating test pretrained script to evaluate model outputs

test_pretrained.ipynb

@@ -16,7 +16,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 83,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
@@ -26,9 +26,9 @@
       "Total dataset examples: 1044\n",
       "\n",
       "\n",
-      "What is the highest number of assists recorded by the Indiana Pacers in a single home game?\n",
-      "SELECT MAX(ast_home)  FROM game  WHERE team_name_home = 'Indiana Pacers';\n",
-      "44.0\n"
+      "List the full names of all teams founded in the 1980s.\n",
+      "SELECT full_name FROM team WHERE year_founded BETWEEN 1980 AND 1989;\n",
+      "Dallas Mavericks, Miami Heat, Minnesota Timberwolves, Orlando Magic, Charlotte Hornets\n"
      ]
     }
    ],
@@ -56,7 +56,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 84,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -80,50 +80,153 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 85,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
-    "input_text = \"\"\"You are an AI assistant that generates SQLite queries for an NBA database based on user questions. The database consists of two tables:\n",
+    "input_text = \"\"\"You are an AI assistant that converts natural language queries into valid SQLite queries.\n",
+    "Database Schema and Explanations\n",
     "\n",
-    "1. `team` - Stores information about NBA teams.\n",
-    "   - `id`: Unique team identifier.\n",
-    "   - `full_name`: Full team name (e.g., \"Los Angeles Lakers\").\n",
-    "   - `abbreviation`: 3-letter team code (e.g., \"LAL\").\n",
-    "   - `city`, `state`: Location of the team.\n",
-    "   - `year_founded`: The year the team was founded.\n",
+    "team Table\n",
+    "Stores information about NBA teams.\n",
+    "CREATE TABLE IF NOT EXISTS \"team\" (\n",
+    "  \"id\" TEXT PRIMARY KEY,      -- Unique identifier for the team\n",
+    "  \"full_name\" TEXT,           -- Full official name of the team (e.g., \"Los Angeles Lakers\")\n",
+    "  \"abbreviation\" TEXT,        -- Shortened team name (e.g., \"LAL\")\n",
+    "  \"nickname\" TEXT,            -- Commonly used nickname for the team (e.g., \"Lakers\")\n",
+    "  \"city\" TEXT,                -- City where the team is based\n",
+    "  \"state\" TEXT,               -- State where the team is located\n",
+    "  \"year_founded\" REAL         -- Year the team was established\n",
+    ");\n",
     "\n",
-    "2. `game` - Stores details of individual games.\n",
-    "   - `game_date`: Date of the game.\n",
-    "   - `team_id_home`, `team_id_away`: Unique IDs of home and away teams.\n",
-    "   - `team_name_home`, `team_name_away`: Full names of the teams.\n",
-    "   - `pts_home`, `pts_away`: Points scored by home and away teams.\n",
-    "   - `wl_home`: \"W\" if the home team won, \"L\" if they lost.\n",
-    "   - `reb_home`, `reb_away`: Total rebounds.\n",
-    "   - `ast_home`, `ast_away`: Total assists.\n",
-    "   - Other statistics include field goals (`fgm_home`, `fg_pct_home`), three-pointers (`fg3m_home`), free throws (`ftm_home`), and turnovers (`tov_home`).\n",
+    "game Table\n",
+    "Contains detailed statistics for each NBA game, including home and away team performance.\n",
+    "CREATE TABLE IF NOT EXISTS \"game\" (\n",
+    "  \"season_id\" TEXT,            -- Season identifier, formatted as \"2YYYY\" (e.g., \"21970\" for the 1970 season)\n",
+    "  \"team_id_home\" TEXT,         -- ID of the home team (matches \"id\" in team table)\n",
+    "  \"team_abbreviation_home\" TEXT, -- Abbreviation of the home team\n",
+    "  \"team_name_home\" TEXT,       -- Full name of the home team\n",
+    "  \"game_id\" TEXT PRIMARY KEY,  -- Unique identifier for the game\n",
+    "  \"game_date\" TIMESTAMP,       -- Date the game was played (YYYY-MM-DD format)\n",
+    "  \"matchup_home\" TEXT,         -- Matchup details including opponent (e.g., \"LAL vs. BOS\")\n",
+    "  \"wl_home\" TEXT,              -- \"W\" if the home team won, \"L\" if they lost\n",
+    "  \"min\" INTEGER,               -- Total minutes played in the game\n",
+    "  \"fgm_home\" REAL,             -- Field goals made by the home team\n",
+    "  \"fga_home\" REAL,             -- Field goals attempted by the home team\n",
+    "  \"fg_pct_home\" REAL,          -- Field goal percentage of the home team\n",
+    "  \"fg3m_home\" REAL,            -- Three-point field goals made by the home team\n",
+    "  \"fg3a_home\" REAL,            -- Three-point attempts by the home team\n",
+    "  \"fg3_pct_home\" REAL,         -- Three-point field goal percentage of the home team\n",
+    "  \"ftm_home\" REAL,             -- Free throws made by the home team\n",
+    "  \"fta_home\" REAL,             -- Free throws attempted by the home team\n",
+    "  \"ft_pct_home\" REAL,          -- Free throw percentage of the home team\n",
+    "  \"oreb_home\" REAL,            -- Offensive rebounds by the home team\n",
+    "  \"dreb_home\" REAL,            -- Defensive rebounds by the home team\n",
+    "  \"reb_home\" REAL,             -- Total rebounds by the home team\n",
+    "  \"ast_home\" REAL,             -- Assists by the home team\n",
+    "  \"stl_home\" REAL,             -- Steals by the home team\n",
+    "  \"blk_home\" REAL,             -- Blocks by the home team\n",
+    "  \"tov_home\" REAL,             -- Turnovers by the home team\n",
+    "  \"pf_home\" REAL,              -- Personal fouls by the home team\n",
+    "  \"pts_home\" REAL,             -- Total points scored by the home team\n",
+    "  \"plus_minus_home\" INTEGER,   -- Plus/minus rating for the home team\n",
+    "  \"video_available_home\" INTEGER, -- Indicates whether video is available (1 = Yes, 0 = No)\n",
+    "  \"team_id_away\" TEXT,         -- ID of the away team\n",
+    "  \"team_abbreviation_away\" TEXT, -- Abbreviation of the away team\n",
+    "  \"team_name_away\" TEXT,       -- Full name of the away team\n",
+    "  \"matchup_away\" TEXT,         -- Matchup details from the away team’s perspective\n",
+    "  \"wl_away\" TEXT,              -- \"W\" if the away team won, \"L\" if they lost\n",
+    "  \"fgm_away\" REAL,             -- Field goals made by the away team\n",
+    "  \"fga_away\" REAL,             -- Field goals attempted by the away team\n",
+    "  \"fg_pct_away\" REAL,          -- Field goal percentage of the away team\n",
+    "  \"fg3m_away\" REAL,            -- Three-point field goals made by the away team\n",
+    "  \"fg3a_away\" REAL,            -- Three-point attempts by the away team\n",
+    "  \"fg3_pct_away\" REAL,         -- Three-point field goal percentage of the away team\n",
+    "  \"ftm_away\" REAL,             -- Free throws made by the away team\n",
+    "  \"fta_away\" REAL,             -- Free throws attempted by the away team\n",
+    "  \"ft_pct_away\" REAL,          -- Free throw percentage of the away team\n",
+    "  \"oreb_away\" REAL,            -- Offensive rebounds by the away team\n",
+    "  \"dreb_away\" REAL,            -- Defensive rebounds by the away team\n",
+    "  \"reb_away\" REAL,             -- Total rebounds by the away team\n",
+    "  \"ast_away\" REAL,             -- Assists by the away team\n",
+    "  \"stl_away\" REAL,             -- Steals by the away team\n",
+    "  \"blk_away\" REAL,             -- Blocks by the away team\n",
+    "  \"tov_away\" REAL,             -- Turnovers by the away team\n",
+    "  \"pf_away\" REAL,              -- Personal fouls by the away team\n",
+    "  \"pts_away\" REAL,             -- Total points scored by the away team\n",
+    "  \"plus_minus_away\" INTEGER,   -- Plus/minus rating for the away team\n",
+    "  \"video_available_away\" INTEGER, -- Indicates whether video is available (1 = Yes, 0 = No)\n",
+    "  \"season_type\" TEXT           -- Regular season or playoffs\n",
+    ");\n",
     "\n",
-    "### Instructions:\n",
-    "- Generate a valid SQLite query to retrieve relevant data from the database.\n",
-    "- Use column names correctly based on the provided schema.\n",
-    "- Ensure the query is well-structured and avoids unnecessary joins.\n",
-    "- Format the query with proper indentation.\n",
+    "other_stats Table\n",
+    "Stores additional game statistics, linked to the game table via game_id.\n",
+    "CREATE TABLE IF NOT EXISTS \"other_stats\" (\n",
+    "  \"game_id\" TEXT,             -- Unique game identifier (links to \"game\" table)\n",
+    "  \"league_id\" TEXT,           -- League identifier\n",
+    "  \"team_id_home\" TEXT,        -- Home team identifier\n",
+    "  \"team_abbreviation_home\" TEXT, -- Home team abbreviation\n",
+    "  \"team_city_home\" TEXT,      -- Home team city\n",
+    "  \"pts_paint_home\" INTEGER,   -- Points in the paint by the home team\n",
+    "  \"pts_2nd_chance_home\" INTEGER, -- Second chance points by the home team\n",
+    "  \"pts_fb_home\" INTEGER,      -- Fast break points by the home team\n",
+    "  \"largest_lead_home\" INTEGER,-- Largest lead by the home team\n",
+    "  \"lead_changes\" INTEGER,     -- Number of lead changes in the game\n",
+    "  \"times_tied\" INTEGER,       -- Number of times the score was tied\n",
+    "  \"team_turnovers_home\" INTEGER, -- Home team turnovers\n",
+    "  \"total_turnovers_home\" INTEGER, -- Total turnovers in the game\n",
+    "  \"team_rebounds_home\" INTEGER, -- Home team rebounds\n",
+    "  \"pts_off_to_home\" INTEGER,  -- Points off turnovers by the home team\n",
+    "  \"team_id_away\" TEXT,        -- Away team identifier\n",
+    "  \"pts_paint_away\" INTEGER,   -- Points in the paint by the away team\n",
+    "  \"pts_2nd_chance_away\" INTEGER, -- Second chance points by the away team\n",
+    "  \"pts_fb_away\" INTEGER,      -- Fast break points by the away team\n",
+    "  \"largest_lead_away\" INTEGER,-- Largest lead by the away team\n",
+    "  \"team_turnovers_away\" INTEGER, -- Away team turnovers\n",
+    "  \"total_turnovers_away\" INTEGER, -- Total turnovers in the game\n",
+    "  \"team_rebounds_away\" INTEGER, -- Away team rebounds\n",
+    "  \"pts_off_to_away\" INTEGER   -- Points off turnovers by the away team\n",
+    ");\n",
     "\n",
-    "### Example Queries:\n",
-    "User: \"What is the most points the Los Angeles Lakers have ever scored at home?\"\n",
+    "\n",
+    "Query Guidelines\n",
+    "Use team_name_home and team_name_away to match teams.\n",
+    "\n",
+    "To filter by season, use season_id = '2YYYY'.\n",
+    "\n",
+    "Example: To get games from 2005, use season_id = '22005'.\n",
+    "\n",
+    "The game_id column links the game and other_stats tables.\n",
+    "\n",
+    "Ensure queries return relevant columns and avoid unnecessary joins.\n",
+    "\n",
+    "Example User Requests and SQLite Queries\n",
+    "Request:\n",
+    "\"What is the most points the Los Angeles Lakers have ever scored at home?\"\n",
     "SQLite:\n",
     "SELECT MAX(pts_home) \n",
     "FROM game \n",
     "WHERE team_name_home = 'Los Angeles Lakers';\n",
     "\n",
-    "User: \"List all games where the Golden State Warriors scored more than 130 points.\" \n",
+    "Request:\n",
+    "\"How many points did the Miami Heat score on January 10, 2010?\"\n",
+    "SQLite:\n",
+    "SELECT team_name_home, pts_home, team_name_away, pts_away \n",
+    "FROM game \n",
+    "WHERE DATE(game_date) = '2010-01-10' \n",
+    "AND (team_name_home = 'Miami Heat' OR team_name_away = 'Miami Heat');\n",
+    "\n",
+    "Request:\n",
+    "\"Which team won the most home games in the 2000 season?\"\n",
     "SQLite:\n",
-    "SELECT game_date, team_name_home, pts_home, team_name_away, pts_away\n",
+    "SELECT team_name_home, COUNT(*) AS wins\n",
     "FROM game\n",
-    "WHERE (team_name_home = 'Golden State Warriors' AND pts_home > 130)\n",
-    "   OR (team_name_away = 'Golden State Warriors' AND pts_away > 130);\n",
-    "   \n",
-    "Now, generate a SQL query based on the following user request: \"\"\""
+    "WHERE wl_home = 'W' AND season_id = '22000'\n",
+    "GROUP BY team_name_home\n",
+    "ORDER BY wins DESC\n",
+    "LIMIT 1;\n",
+    "\n",
+    "Generate only the SQLite query prefaced by SQLite: and no other text, do not output an explanation of the query. Now generate an SQLite query for the following question: \"\"\""
    ]
   },
   {
@@ -135,7 +238,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 86,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [
     {
@@ -153,9 +256,9 @@
      "output_type": "stream",
      "text": [
       "SQLite:\n",
-      "SELECT MAX(ast_home) \n",
-      "FROM game \n",
-      "WHERE team_name_home = 'Indiana Pacers';\n",
+      "SELECT full_name \n",
+      "FROM team \n",
+      "WHERE year_founded BETWEEN 1980 AND 1989;\n",
       "\n"
      ]
     }
@@ -180,7 +283,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
     {
@@ -188,7 +291,11 @@
      "output_type": "stream",
      "text": [
       "cleaned\n",
-      "(44.0,)\n"
+      "('Dallas Mavericks',)\n",
+      "('Miami Heat',)\n",
+      "('Minnesota Timberwolves',)\n",
+      "('Orlando Magic',)\n",
+      "('Charlotte Hornets',)\n"
      ]
     }
    ],
@@ -222,25 +329,34 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 67,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
+      "Setting `pad_token_id` to `eos_token_id`:32021 for open-end generation.\n"
+     ]
+    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "cleaned\n",
-      "[(44.0,)]\n",
-      "\n",
-      "SELECT MAX(ast_home) \n",
+      "How many times did the Minnesota Timberwolves lose at home in the 2004 season despite recording more steals and blocks than their opponent?\n",
+      "SELECT COUNT(*) FROM game g WHERE g.team_abbreviation_home = 'MIN' AND g.wl_home = 'L' AND g.stl_home > g.stl_away AND g.blk_home > g.blk_away AND g.season_id = '22004';\n",
+      "0\n",
+      "SQLite:\n",
+      "SELECT COUNT(*) \n",
       "FROM game \n",
-      "WHERE team_name_home = 'Indiana Pacers';\n",
+      "WHERE team_name_home = 'Minnesota Timberwolves' \n",
+      "AND wl_home = 'L' \n",
+      "AND season_id = '22004';\n",
       "\n",
-      "SELECT MAX(ast_home)  FROM game  WHERE team_name_home = 'Indiana Pacers';\n",
-      "44.0\n",
-      "44.0\n",
-      "SQL matched? True\n",
-      "Result matched? True\n"
+      "[(17,)]\n",
+      "SQL matched? False\n",
+      "Result matched? False\n"
      ]
     }
    ],
@@ -260,19 +376,51 @@
     "        cursor.execute(query)\n",
     "        rows = cursor.fetchall()\n",
     "\n",
+    "        # Strip all whitespace before comparing queries since there may be differences in spacing, newlines, tabs, etc.\n",
+    "        query = query.replace(\" \", \"\").replace(\"\\n\", \"\").replace(\"\\t\", \"\")\n",
+    "        sample_query = sample_query.replace(\" \", \"\").replace(\"\\n\", \"\").replace(\"\\t\", \"\")\n",
+    "        query_match = (query == sample_query)\n",
+    "\n",
     "        # Check if this is a multi-line query\n",
     "        if \"|\" in sample_result:\n",
-    "            return True, True\n",
+    "            result_list = sample_result.split(\"|\") \n",
+    "            for i in range(len(result_list)):\n",
+    "                result_list[i] = str(result_list[i]).strip()\n",
+    "            result = False\n",
+    "            for row in rows:\n",
+    "                for r in row:\n",
+    "                    if str(r) in result_list:\n",
+    "                        return query_match, True\n",
+    "            print(rows)\n",
+    "            return query_match, result\n",
     "        else:\n",
-    "            # Strip all whitespace before comparing queries since there may be differences in spacing, newlines, tabs, etc.\n",
-    "            query = query.replace(\" \", \"\").replace(\"\\n\", \"\").replace(\"\\t\", \"\")\n",
-    "            sample_query = sample_query.replace(\" \", \"\").replace(\"\\n\", \"\").replace(\"\\t\", \"\")\n",
+    "            print(rows)\n",
+    "            result = False\n",
+    "            for row in rows:\n",
+    "                for r in row:\n",
+    "                    if str(r) == str(sample_result):\n",
+    "                        return query_match, True\n",
     "\n",
     "            # Compare results and return\n",
-    "            return (query == sample_query), (str(rows[0][0]) == str(sample_result))\n",
+    "            return query_match, result\n",
     "    except:\n",
     "        return False, False\n",
     "\n",
+    "# Obtain sample\n",
+    "sample = df.sample(n=1)\n",
+    "print(sample[\"natural_query\"].values[0])\n",
+    "print(sample[\"sql_query\"].values[0])\n",
+    "print(sample[\"result\"].values[0])\n",
+    "\n",
+    "# Create message with sample query and run model\n",
+    "message=[{ 'role': 'user', 'content': input_text + sample[\"natural_query\"].values[0]}]\n",
+    "inputs = tokenizer.apply_chat_template(message, add_generation_prompt=True, return_tensors=\"pt\").to(model.device)\n",
+    "outputs = model.generate(inputs, max_new_tokens=512, do_sample=False, top_k=50, top_p=0.95, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)\n",
+    "\n",
+    "# Print output\n",
+    "query_output = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)\n",
+    "print(query_output)\n",
+    "\n",
     "result = compare_result(sample[\"sql_query\"].values[0], sample[\"result\"].values[0], query_output)\n",
     "print(\"SQL matched? \" + str(result[0]))\n",
     "print(\"Result matched? \" + str(result[1]))"

train-data/sql_train.tsv

@@ -470,7 +470,7 @@ What is the highest combined pts in any game involving the Miami Heat?	SELECT MA
 How many away games did the Milwaukee Bucks play in the 2018 season?	SELECT COUNT(*) FROM game WHERE team_name_away = 'Milwaukee Bucks' AND season_id = '22018';	41.0
 In which season did the Cleveland Cavaliers have the highest average ft_pct at home?	SELECT season_id, AVG(ft_pct_home) as avg_stat FROM game WHERE team_name_home = 'Cleveland Cavaliers' GROUP BY season_id ORDER BY avg_stat DESC LIMIT 1;	2014.0
 In which season did the Golden State Warriors have the highest average ft_pct at home?	SELECT season_id, AVG(ft_pct_home) as avg_stat FROM game WHERE team_name_home = 'Golden State Warriors' GROUP BY season_id ORDER BY avg_stat DESC LIMIT 1;	2016.0
-In which season did the Houston Rockets have the highest average reb at home?	SELECT season_id, AVG(reb_home) as avg_stat FROM game WHERE team_name_home = 'Houston Rockets' GROUP BY season_id ORDER BY avg_stat DESC LIMIT 1;	1984.0
+In which season did the Houston Rockets have the highest average reb at home?	SELECT season_id, AVG(reb_home) as avg_stat FROM game WHERE team_name_home = 'Houston Rockets' GROUP BY season_id ORDER BY avg_stat DESC LIMIT 1;	41984|58.0
 What is the highest combined ast in any game involving the Los Angeles Lakers?	SELECT MAX(ast_home + ast_away) FROM game WHERE team_name_home = 'Los Angeles Lakers' OR team_name_away = 'Los Angeles Lakers';	86.0
 How many away games did the Chicago Bulls play in the 2022 season?	SELECT COUNT(*) FROM game WHERE team_name_away = 'Chicago Bulls' AND season_id = '22022';	41.0
 How many home games did the Boston Celtics play in the 2018 season?	SELECT COUNT(*) FROM game WHERE team_name_home = 'Boston Celtics' AND season_id = '22018';	41.0