Updating compare result function in src directory, fixing small issues in test_rag prompt, adding test_rag notebook

src/evaluation/compare_result.py

@@ -1,13 +1,27 @@
 import math
 
-def compare_result(cursor, sample_query, sample_result, query_output):
+def compare_result(sample_query, sample_result, query_output):
     # Clean model output to only have the query output
-    if query_output[0:7] == "SQLite:":
+    if query_output[0:8] == "SQLite:\n":
+        query = query_output[8:]
+    elif query_output[0:8] == "SQLite: ":
+        query = query_output[8:]
+    elif query_output[0:7] == "SQLite:":
         query = query_output[7:]
+    elif query_output[0:5] == "SQL:\n":
+        query = query_output[5:]
+    elif query_output[0:5] == "SQL: ":
+        query = query_output[5:]
     elif query_output[0:4] == "SQL:":
         query = query_output[4:]
     else:
         query = query_output
+
+    # Clean any excess text after the query semicolon
+    for i in range(len(query)):
+        if query[i] == ";":
+            query = query[:i+1]
+            break
     
     # Try to execute query, if it fails, then this is a failure of the model
     try:
@@ -47,7 +61,7 @@ def compare_result(cursor, sample_query, sample_result, query_output):
                             if math.isclose(float(r), float(res), abs_tol=0.5):
                                 return True, query_match, True
                         except:
-                            if r in res or res in r:
+                            if str(r) in res or res in str(r):
                                 return True, query_match, True
                     
             # Check if the model returned a sum of examples as opposed to the whole thing
@@ -90,4 +104,4 @@ def compare_result(cursor, sample_query, sample_result, query_output):
             # Compare results and return
             return True, query_match, result
     except:
-        return False, False, False
+        return False, False, False

test_rag.ipynb

@@ -0,0 +1,262 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "9ba5b9ac",
+   "metadata": {},
+   "source": [
+    "# Notebook to evaluate RAG performance"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b7c75665",
+   "metadata": {},
+   "source": [
+    "## Create RAG document store"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "d589714b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Users\\Dean\\AppData\\Local\\Programs\\Python\\Python312\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "WARNING:tensorflow:From c:\\Users\\Dean\\AppData\\Local\\Programs\\Python\\Python312\\Lib\\site-packages\\tf_keras\\src\\losses.py:2976: The name tf.losses.sparse_softmax_cross_entropy is deprecated. Please use tf.compat.v1.losses.sparse_softmax_cross_entropy instead.\n",
+      "\n",
+      "Total dataset examples: 1044\n",
+      "\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "import pandas as pd\n",
+    "import warnings\n",
+    "import torch\n",
+    "import time\n",
+    "import math\n",
+    "import sqlite3 as sql\n",
+    "\n",
+    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
+    "from rag_metadata import SQLMetadataRetriever\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\")\n",
+    "\n",
+    "# Establish a database connection once (adjust the DB path as needed)\n",
+    "connection = sql.connect('./nba-data/nba.sqlite')\n",
+    "cursor = connection.cursor()\n",
+    "\n",
+    "# ------------------------------\n",
+    "# Load dataset and print summary\n",
+    "# ------------------------------\n",
+    "df = pd.read_csv(\"./train-data/sql_train.tsv\", sep='\\t')\n",
+    "print(\"Total dataset examples: \" + str(len(df)))\n",
+    "print(\"\\n\")\n",
+    "\n",
+    "# ------------------------------\n",
+    "# Load tokenizer and model\n",
+    "# ------------------------------\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\"./deepseek-coder-1.3b-instruct\")\n",
+    "model = AutoModelForCausalLM.from_pretrained(\n",
+    "    \"./deepseek-coder-1.3b-instruct\",\n",
+    "    torch_dtype=torch.bfloat16,\n",
+    "    device_map=device\n",
+    ")\n",
+    "model.generation_config.pad_token_id = tokenizer.pad_token_id\n",
+    "\n",
+    "# ------------------------------\n",
+    "# Initialize RAG retriever and load schema metadata\n",
+    "# ------------------------------\n",
+    "retriever = SQLMetadataRetriever()\n",
+    "\n",
+    "metadata_docs = [\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",
+    "    '''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",
+    "    '''other_stats Table\n",
+    "Stores additional statistics, linked to the game table via game_id.\n",
+    "CREATE TABLE IF NOT EXISTS \"other_stats\" (\n",
+    "  \"game_id\" TEXT,             -- Unique game identifier, matches id column from 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 \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 by the home team\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",
+    "  \"team_abbreviation_away\" TEXT,  -- Away team abbreviation\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 by the away team\n",
+    "  \"team_rebounds_away\" INTEGER, -- Away team rebounds\n",
+    "  \"pts_off_to_away\" INTEGER   -- Points off turnovers by the away team\n",
+    ");\n",
+    "''',\n",
+    "    '''Team Name Information\n",
+    "In plaintext user questions, only the full team names will be used, but in the queries you may use either full names or abbreviations.\n",
+    "Full names are used with the game table, while abbreviations should be used with the other_stats table.\n",
+    "Team names and abbreviations (separated by |):\n",
+    "Atlanta Hawks|ATL, Boston Celtics|BOS, Cleveland Cavaliers|CLE, New Orleans Pelicans|NOP,\n",
+    "Chicago Bulls|CHI, Dallas Mavericks|DAL, Denver Nuggets|DEN, Golden State Warriors|GSW,\n",
+    "Houston Rockets|HOU, Los Angeles Clippers|LAC, Los Angeles Lakers|LAL, Miami Heat|MIA,\n",
+    "Milwaukee Bucks|MIL, Minnesota Timberwolves|MIN, Brooklyn Nets|BKN, New York Knicks|NYK,\n",
+    "Orlando Magic|ORL, Indiana Pacers|IND, Philadelphia 76ers|PHI, Phoenix Suns|PHX,\n",
+    "Portland Trail Blazers|POR, Sacramento Kings|SAC, San Antonio Spurs|SAS,\n",
+    "Oklahoma City Thunder|OKC, Toronto Raptors|TOR, Utah Jazz|UTA, Memphis Grizzlies|MEM,\n",
+    "Washington Wizards|WAS, Detroit Pistons|DET, Charlotte Hornets|CHA\n",
+    "'''\n",
+    "]\n",
+    "\n",
+    "retriever.add_documents(metadata_docs)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "499d2745",
+   "metadata": {},
+   "source": [
+    "## Define compare result function for evaluation process"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "268561cd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e7393ccb",
+   "metadata": {},
+   "source": [
+    "## Evaluate RAG model on single training example"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "500f003b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.12.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

test_rag.py

@@ -282,15 +282,15 @@ def run_evaluation(nba_df, title):
         
         # Build the prompt with instructions, schema, examples, and current request.
         input_text = f"""
-You are an AI assistant that generates SQL queries for an NBA database based on user questions.
+You are an AI assistant that generates SQLite queries for an NBA database based on user questions.
 
 ### Relevant Schema:
 {schema_block}
 
 ### Instructions:
-- Generate a valid SQL query to retrieve relevant data from the database.
+- Generate a valid SQLite query to retrieve relevant data from the database.
 - Use column names correctly based on the provided schema.
-- Output only the SQL query as plain text.
+- Output only the SQLite query as plain text.
 
 ### Example Queries:
 Use team_name_home and team_name_away to match teams to the game table.