Add a source folder

src/evaluation/compare_result.py

@@ -0,0 +1,93 @@
+import math
+
+def compare_result(cursor, sample_query, sample_result, query_output):
+    # Clean model output to only have the query output
+    if query_output[0:7] == "SQLite:":
+        query = query_output[7:]
+    elif query_output[0:4] == "SQL:":
+        query = query_output[4:]
+    else:
+        query = query_output
+    
+    # Try to execute query, if it fails, then this is a failure of the model
+    try:
+        # Execute query and obtain result
+        cursor.execute(query)
+        rows = cursor.fetchall()
+
+        # Strip all whitespace before comparing queries since there may be differences in spacing, newlines, tabs, etc.
+        query = query.replace(" ", "").replace("\n", "").replace("\t", "")
+        sample_query = sample_query.replace(" ", "").replace("\n", "").replace("\t", "")
+        query_match = (query == sample_query)
+
+        # If the queries match, the results clearly also match
+        if query_match:
+            return True, True, True
+
+        # Check if this is a multi-line query
+        if "|" in sample_result or "(" in sample_result:
+            #print(rows)
+            # Create list of results by stripping separators and splitting on them
+            if "(" in sample_result:
+                sample_result = sample_result.replace("(", "").replace(")", "")
+                result_list = sample_result.split(",") 
+            else:
+                result_list = sample_result.split("|") 
+
+            # Strip all results in list
+            for i in range(len(result_list)):
+                result_list[i] = str(result_list[i]).strip()
+            
+            # Loop through model result and see if it matches training example
+            result = False
+            for row in rows:
+                for r in row:
+                    for res in result_list:
+                        try:
+                            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:
+                                return True, query_match, True
+                    
+            # Check if the model returned a sum of examples as opposed to the whole thing
+            if len(rows) == 1:
+                for r in rows[0]:
+                    if r == str(len(result_list)):
+                        return True, query_match, True
+                    
+            return True, query_match, result
+        # Else the sample result is a single value or string
+        else:
+            #print(rows)
+            result = False
+            # Loop through model result and see if it contains the sample result
+            for row in rows:
+                for r in row:
+                    # Check by string
+                    if str(r) in str(sample_result):
+                        try:
+                            if math.isclose(float(r), float(sample_result), abs_tol=0.5):
+                                return True, query_match, True
+                        except:
+                            return True, query_match, True
+                    # Check by number, using try incase the cast as float fails
+                    try:
+                        if math.isclose(float(r), float(sample_result), abs_tol=0.5):
+                            return True, query_match, True
+                    except:
+                        pass
+
+            # Check if the model returned a list of examples instead of a total sum (both acceptable)
+            try:
+                if len(rows) > 1 and len(rows) == int(sample_result):
+                    return True, query_match, True
+                if len(rows[0]) > 1 and rows[0][1] is not None and  len(rows[0]) == int(sample_result):
+                    return True, query_match, True
+            except:
+                pass
+
+            # Compare results and return
+            return True, query_match, result
+    except:
+        return False, False, False

src/prompts/prompt.py

@@ -0,0 +1,179 @@
+input_text = """You are an AI assistant that converts natural language queries into valid SQLite queries.
+Database Schema and Explanations
+
+team Table
+Stores information about NBA teams.
+CREATE TABLE IF NOT EXISTS "team" (
+  "id" TEXT PRIMARY KEY,      -- Unique identifier for the team
+  "full_name" TEXT,           -- Full official name of the team (e.g., "Los Angeles Lakers")
+  "abbreviation" TEXT,        -- Shortened team name (e.g., "LAL")
+  "nickname" TEXT,            -- Commonly used nickname for the team (e.g., "Lakers")
+  "city" TEXT,                -- City where the team is based
+  "state" TEXT,               -- State where the team is located
+  "year_founded" REAL         -- Year the team was established
+);
+
+game Table
+Contains detailed statistics for each NBA game, including home and away team performance.
+CREATE TABLE IF NOT EXISTS "game" (
+  "season_id" TEXT,            -- Season identifier, formatted as "2YYYY" (e.g., "21970" for the 1970 season)
+  "team_id_home" TEXT,         -- ID of the home team (matches "id" in team table)
+  "team_abbreviation_home" TEXT, -- Abbreviation of the home team
+  "team_name_home" TEXT,       -- Full name of the home team
+  "game_id" TEXT PRIMARY KEY,  -- Unique identifier for the game
+  "game_date" TIMESTAMP,       -- Date the game was played (YYYY-MM-DD format)
+  "matchup_home" TEXT,         -- Matchup details including opponent (e.g., "LAL vs. BOS")
+  "wl_home" TEXT,              -- "W" if the home team won, "L" if they lost
+  "min" INTEGER,               -- Total minutes played in the game
+  "fgm_home" REAL,             -- Field goals made by the home team
+  "fga_home" REAL,             -- Field goals attempted by the home team
+  "fg_pct_home" REAL,          -- Field goal percentage of the home team
+  "fg3m_home" REAL,            -- Three-point field goals made by the home team
+  "fg3a_home" REAL,            -- Three-point attempts by the home team
+  "fg3_pct_home" REAL,         -- Three-point field goal percentage of the home team
+  "ftm_home" REAL,             -- Free throws made by the home team
+  "fta_home" REAL,             -- Free throws attempted by the home team
+  "ft_pct_home" REAL,          -- Free throw percentage of the home team
+  "oreb_home" REAL,            -- Offensive rebounds by the home team
+  "dreb_home" REAL,            -- Defensive rebounds by the home team
+  "reb_home" REAL,             -- Total rebounds by the home team
+  "ast_home" REAL,             -- Assists by the home team
+  "stl_home" REAL,             -- Steals by the home team
+  "blk_home" REAL,             -- Blocks by the home team
+  "tov_home" REAL,             -- Turnovers by the home team
+  "pf_home" REAL,              -- Personal fouls by the home team
+  "pts_home" REAL,             -- Total points scored by the home team
+  "plus_minus_home" INTEGER,   -- Plus/minus rating for the home team
+  "video_available_home" INTEGER, -- Indicates whether video is available (1 = Yes, 0 = No)
+  "team_id_away" TEXT,         -- ID of the away team
+  "team_abbreviation_away" TEXT, -- Abbreviation of the away team
+  "team_name_away" TEXT,       -- Full name of the away team
+  "matchup_away" TEXT,         -- Matchup details from the away team’s perspective
+  "wl_away" TEXT,              -- "W" if the away team won, "L" if they lost
+  "fgm_away" REAL,             -- Field goals made by the away team
+  "fga_away" REAL,             -- Field goals attempted by the away team
+  "fg_pct_away" REAL,          -- Field goal percentage of the away team
+  "fg3m_away" REAL,            -- Three-point field goals made by the away team
+  "fg3a_away" REAL,            -- Three-point attempts by the away team
+  "fg3_pct_away" REAL,         -- Three-point field goal percentage of the away team
+  "ftm_away" REAL,             -- Free throws made by the away team
+  "fta_away" REAL,             -- Free throws attempted by the away team
+  "ft_pct_away" REAL,          -- Free throw percentage of the away team
+  "oreb_away" REAL,            -- Offensive rebounds by the away team
+  "dreb_away" REAL,            -- Defensive rebounds by the away team
+  "reb_away" REAL,             -- Total rebounds by the away team
+  "ast_away" REAL,             -- Assists by the away team
+  "stl_away" REAL,             -- Steals by the away team
+  "blk_away" REAL,             -- Blocks by the away team
+  "tov_away" REAL,             -- Turnovers by the away team
+  "pf_away" REAL,              -- Personal fouls by the away team
+  "pts_away" REAL,             -- Total points scored by the away team
+  "plus_minus_away" INTEGER,   -- Plus/minus rating for the away team
+  "video_available_away" INTEGER, -- Indicates whether video is available (1 = Yes, 0 = No)
+  "season_type" TEXT           -- Regular season or playoffs
+);
+
+other_stats Table
+Stores additional statistics, linked to the game table via game_id.
+CREATE TABLE IF NOT EXISTS "other_stats" (
+  "game_id" TEXT,             -- Unique game identifier, matches id column from game table
+  "league_id" TEXT,           -- League identifier
+  "team_id_home" TEXT,        -- Home team identifier
+  "team_abbreviation_home" TEXT, -- Home team abbreviation
+  "team_city_home" TEXT,      -- Home team city
+  "pts_paint_home" INTEGER,   -- Points in the paint by the home team
+  "pts_2nd_chance_home" INTEGER, -- Second chance points by the home team
+  "pts_fb_home" INTEGER,      -- Fast break points by the home team
+  "largest_lead_home" INTEGER,-- Largest lead by the home team
+  "lead_changes" INTEGER,     -- Number of lead changes 
+  "times_tied" INTEGER,       -- Number of times the score was tied
+  "team_turnovers_home" INTEGER, -- Home team turnovers
+  "total_turnovers_home" INTEGER, -- Total turnovers by the home team
+  "team_rebounds_home" INTEGER, -- Home team rebounds
+  "pts_off_to_home" INTEGER,  -- Points off turnovers by the home team
+  "team_id_away" TEXT,        -- Away team identifier
+  "team_abbreviation_away" TEXT,  -- Away team abbreviation
+  "pts_paint_away" INTEGER,   -- Points in the paint by the away team
+  "pts_2nd_chance_away" INTEGER, -- Second chance points by the away team
+  "pts_fb_away" INTEGER,      -- Fast break points by the away team
+  "largest_lead_away" INTEGER,-- Largest lead by the away team
+  "team_turnovers_away" INTEGER, -- Away team turnovers
+  "total_turnovers_away" INTEGER, -- Total turnovers by the away team
+  "team_rebounds_away" INTEGER, -- Away team rebounds
+  "pts_off_to_away" INTEGER   -- Points off turnovers by the away team
+);
+
+
+Team Name Information
+In the plaintext user questions, only the full team names will be used, but in the queries you may use the full team names or the abbreviations. 
+The full team names can be used with the game table, while the abbreviations should be used with the other_stats table.
+Notice they are separated by the | character in the following list:
+
+Atlanta Hawks|ATL
+Boston Celtics|BOS
+Cleveland Cavaliers|CLE
+New Orleans Pelicans|NOP
+Chicago Bulls|CHI
+Dallas Mavericks|DAL
+Denver Nuggets|DEN
+Golden State Warriors|GSW
+Houston Rockets|HOU
+Los Angeles Clippers|LAC
+Los Angeles Lakers|LAL
+Miami Heat|MIA
+Milwaukee Bucks|MIL
+Minnesota Timberwolves|MIN
+Brooklyn Nets|BKN
+New York Knicks|NYK
+Orlando Magic|ORL
+Indiana Pacers|IND
+Philadelphia 76ers|PHI
+Phoenix Suns|PHX
+Portland Trail Blazers|POR
+Sacramento Kings|SAC
+San Antonio Spurs|SAS
+Oklahoma City Thunder|OKC
+Toronto Raptors|TOR
+Utah Jazz|UTA
+Memphis Grizzlies|MEM
+Washington Wizards|WAS
+Detroit Pistons|DET
+Charlotte Hornets|CHA
+
+Query Guidelines
+Use team_name_home and team_name_away to match teams to the game table. Use team_abbreviation_home and team_abbreviation away to match teams to the other_stats table.
+
+To filter by season, use season_id = '2YYYY'.
+
+Example: To get statistics from 2005, use a statement like: season_id = '22005'. To get statistics from 1972, use a statement like: season_id = "21972". To get statistics from 2015, use a statement like: season_id = "22015".
+
+Ensure queries return relevant columns and avoid unnecessary joins.
+
+Example User Requests and SQLite Queries
+Request:
+"What is the most points the Los Angeles Lakers have ever scored at home?"
+SQLite:
+SELECT MAX(pts_home) FROM game WHERE team_name_home = 'Los Angeles Lakers';
+
+Request:
+"Which teams are located in the state of California?"
+SQLite:
+SELECT full_name FROM team WHERE state = 'California';
+
+Request:
+"Which team had the highest number of team turnovers in an away game?"
+SQLite:
+SELECT team_abbreviation_away FROM other_stats ORDER BY team_turnovers_away DESC LIMIT 1;
+
+Request:
+"Which teams were founded before 1979?"
+SQLite:
+SELECT full_name FROM team WHERE year_founded < 1979;
+
+Request:
+"Find the Boston Celtics largest home victory margin in the 2008 season."
+SQLite:
+SELECT MAX(pts_home - pts_away) AS biggest_win FROM game WHERE team_name_home = 'Boston Celtics' AND season_id = '22008';
+
+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 user request. Request:
+"""

test_pretrained.ipynb

@@ -16,7 +16,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -26,9 +26,9 @@
       "Total dataset examples: 1044\n",
       "\n",
       "\n",
-      "Which team committed the fewest total turnovers in an away game that resulted in a win?\n",
-      "SELECT team_abbreviation_away FROM other_stats WHERE game_id IN (SELECT game_id FROM game WHERE wl_away = 'W') ORDER BY total_turnovers_away ASC LIMIT 1;\n",
-      "PHX\n"
+      "Which team had the largest lead in a single game in the 2001 season?\n",
+      "SELECT g.team_name_home AS team, os.largest_lead_home AS lead FROM other_stats os JOIN game g ON os.game_id = g.game_id WHERE g.season_id = '22001' ORDER BY os.largest_lead_home DESC LIMIT 1;\n",
+      "Portland Trail Blazers|47\n"
      ]
     }
    ],
@@ -58,7 +58,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -83,193 +83,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [],
    "source": [
-    "input_text = \"\"\"You are an AI assistant that converts natural language queries into valid SQLite queries.\n",
-    "Database Schema and Explanations\n",
-    "\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",
-    "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",
-    "\n",
-    "Team Name Information\n",
-    "In the plaintext user questions, only the full team names will be used, but in the queries you may use the full team names or the abbreviations. \n",
-    "The full team names can be used with the game table, while the abbreviations should be used with the other_stats table.\n",
-    "Notice they are separated by the | character in the following list:\n",
-    "\n",
-    "Atlanta Hawks|ATL\n",
-    "Boston Celtics|BOS\n",
-    "Cleveland Cavaliers|CLE\n",
-    "New Orleans Pelicans|NOP\n",
-    "Chicago Bulls|CHI\n",
-    "Dallas Mavericks|DAL\n",
-    "Denver Nuggets|DEN\n",
-    "Golden State Warriors|GSW\n",
-    "Houston Rockets|HOU\n",
-    "Los Angeles Clippers|LAC\n",
-    "Los Angeles Lakers|LAL\n",
-    "Miami Heat|MIA\n",
-    "Milwaukee Bucks|MIL\n",
-    "Minnesota Timberwolves|MIN\n",
-    "Brooklyn Nets|BKN\n",
-    "New York Knicks|NYK\n",
-    "Orlando Magic|ORL\n",
-    "Indiana Pacers|IND\n",
-    "Philadelphia 76ers|PHI\n",
-    "Phoenix Suns|PHX\n",
-    "Portland Trail Blazers|POR\n",
-    "Sacramento Kings|SAC\n",
-    "San Antonio Spurs|SAS\n",
-    "Oklahoma City Thunder|OKC\n",
-    "Toronto Raptors|TOR\n",
-    "Utah Jazz|UTA\n",
-    "Memphis Grizzlies|MEM\n",
-    "Washington Wizards|WAS\n",
-    "Detroit Pistons|DET\n",
-    "Charlotte Hornets|CHA\n",
-    "\n",
-    "Query Guidelines\n",
-    "Use team_name_home and team_name_away to match teams to the game table. Use team_abbreviation_home and team_abbreviation away to match teams to the other_stats table.\n",
-    "\n",
-    "To filter by season, use season_id = '2YYYY'.\n",
-    "\n",
-    "Example: To get statistics from 2005, use a statement like: season_id = '22005'. To get statistics from 1972, use a statement like: season_id = \"21972\". To get statistics from 2015, use a statement like: season_id = \"22015\".\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",
-    "Request:\n",
-    "\"Which teams are located in the state of California?\"\n",
-    "SQLite:\n",
-    "SELECT full_name FROM team WHERE state = 'California';\n",
-    "\n",
-    "Request:\n",
-    "\"Which team had the highest number of team turnovers in an away game?\"\n",
-    "SQLite:\n",
-    "SELECT team_abbreviation_away FROM other_stats ORDER BY team_turnovers_away DESC LIMIT 1;\n",
-    "\n",
-    "Request:\n",
-    "\"Which teams were founded before 1979?\"\n",
-    "SQLite:\n",
-    "SELECT full_name FROM team WHERE year_founded < 1979;\n",
-    "\n",
-    "Request:\n",
-    "\"Find the Boston Celtics largest home victory margin in the 2008 season.\"\n",
-    "SQLite:\n",
-    "SELECT MAX(pts_home - pts_away) AS biggest_win\n",
-    "FROM game\n",
-    "WHERE team_name_home = 'Boston Celtics' AND season_id = '22008';\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 user request. Request:\n",
-    "\"\"\""
+    "from src.prompts.prompt import input_text"
    ]
   },
   {
@@ -281,7 +99,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
@@ -289,9 +107,7 @@
      "output_type": "stream",
      "text": [
       "SQLite:\n",
-      "SELECT team_abbreviation_away \n",
-      "FROM other_stats \n",
-      "WHERE wl_away = 'W' AND total_turnovers_away < (SELECT MIN(total_turnovers_away) FROM other_stats WHERE wl_away = 'L');\n",
+      "SELECT team_abbreviation_home FROM other_stats WHERE lead_changes = 1 AND season_id = '2001';\n",
       "\n"
      ]
     }
@@ -316,7 +132,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [
     {
@@ -361,29 +177,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "What is the largest lead the Minnesota Timberwolves had at home?\n",
-      "SELECT MAX(largest_lead_home) as max_lead  FROM other_stats  WHERE team_abbreviation_home = 'MIN';\n",
-      "48.0\n",
-      "SQLite:\n",
-      "SELECT MAX(largest_lead_home) \n",
-      "FROM other_stats \n",
-      "WHERE team_name_home = 'Minnesota Timberwolves';\n",
-      "\n",
-      "Statement valid? False\n",
-      "SQLite matched? False\n",
-      "Result matched? False\n"
+     "ename": "ImportError",
+     "evalue": "cannot import name 'compare_result_two' from 'src.evaluation.compare_result' (/Users/esteban/Documents/USC/spring_2025/NLP/SQL-Generation/src/evaluation/compare_result.py)",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[30], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmath\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mevaluation\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcompare_result\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m compare_result_two\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcompare_result\u001b[39m(sample_query, sample_result, query_output):\n\u001b[1;32m      5\u001b[0m     \u001b[38;5;66;03m# Clean model output to only have the query output\u001b[39;00m\n\u001b[1;32m      6\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m query_output[\u001b[38;5;241m0\u001b[39m:\u001b[38;5;241m7\u001b[39m] \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mSQLite:\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n",
+      "\u001b[0;31mImportError\u001b[0m: cannot import name 'compare_result_two' from 'src.evaluation.compare_result' (/Users/esteban/Documents/USC/spring_2025/NLP/SQL-Generation/src/evaluation/compare_result.py)"
      ]
     }
    ],
    "source": [
     "import math\n",
+    "from src.evaluation.compare_result import compare_result_two\n",
     "\n",
     "def compare_result(sample_query, sample_result, query_output):\n",
     "    # Clean model output to only have the query output\n",
@@ -479,9 +290,18 @@
     "\n",
     "# Obtain sample\n",
     "sample = df.sample(n=1)\n",
+    "sample_dic = {\n",
+    "    \"natural_query\": \"How many home games did the Miami Heat play in the 2021 season?\",\n",
+    "    \"sql_query\": \"SELECT COUNT(*) FROM game WHERE team_name_home = 'Miami Heat' AND season_id = '22021';\",\n",
+    "    \"result\": 41.0\n",
+    "}\n",
+    "\n",
+    "sample = pd.DataFrame([sample_dic])\n",
+    "\"\"\"\n",
     "print(sample[\"natural_query\"].values[0])\n",
     "print(sample[\"sql_query\"].values[0])\n",
     "print(sample[\"result\"].values[0])\n",
+    "\"\"\"\n",
     "\n",
     "# Create message with sample query and run model\n",
     "message=[{ 'role': 'user', 'content': input_text + sample[\"natural_query\"].values[0]}]\n",
@@ -495,7 +315,12 @@
     "result = compare_result(sample[\"sql_query\"].values[0], sample[\"result\"].values[0], query_output)\n",
     "print(\"Statement valid? \" + str(result[0]))\n",
     "print(\"SQLite matched? \" + str(result[1]))\n",
-    "print(\"Result matched? \" + str(result[2]))"
+    "print(\"Result matched? \" + str(result[2]))\n",
+    "\n",
+    "result_two = compare_result_two(cursor, sample[\"sql_query\"].values[0], sample[\"result\"].values[0], query_output)\n",
+    "print(\"Statement valid? \" + str(result_two[0]))\n",
+    "print(\"SQLite matched? \" + str(result_two[1]))\n",
+    "print(\"Result matched? \" + str(result_two[2]))"
    ]
   },
   {
@@ -830,7 +655,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "CSCI544",
    "language": "python",
    "name": "python3"
   },
@@ -844,7 +669,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.6"
+   "version": "3.11.11"
   }
  },
  "nbformat": 4,