dvwn commited on
Commit
97dafec
·
1 Parent(s): 9ef286c

NL2SQL Qwen Model V1.0.0

Browse files
app/main.py ADDED
File without changes
evaluation_results.json ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [
2
+ {
3
+ "id": 1,
4
+ "status": "ERROR: Execution failed on sql 'SELECT t.Genre, COUNT(t.TrackId) AS TrackCount FROM Track t GROUP BY t.Genre;': no such column: t.Genre"
5
+ },
6
+ {
7
+ "id": 2,
8
+ "status": "PASS",
9
+ "generated": "SELECT t1.Title, ar.Name FROM Album AS t1 INNER JOIN Artist AS ar ON t1.ArtistId = ar.ArtistId;"
10
+ },
11
+ {
12
+ "id": 3,
13
+ "status": "ERROR: Execution failed on sql 'SELECT T4.Country, SUM(T3.UnitPrice * T3.Quantity) as TotalRevenue\nFROM InvoiceLine AS T1\nJOIN Track AS T2 ON T1.TrackId = T2.TrackId\nJOIN Album AS T3 ON T2.AlbumId = T3.AlbumId\nJOIN Genre AS T4 ON T2.GenreId = T4.GenreId\nGROUP BY T4.Country;': no such column: T4.Country"
14
+ },
15
+ {
16
+ "id": 4,
17
+ "status": "ERROR: Execution failed on sql 'SELECT e.FirstName, e.LastName FROM Employee AS e JOIN Customer AS c ON e.SupportRepId = c.CustomerId WHERE c.Company LIKE 'Sales Support';': no such column: e.SupportRepId"
18
+ },
19
+ {
20
+ "id": 5,
21
+ "status": "FAIL",
22
+ "generated": "SELECT c.FirstName, c.LastName, SUM(il.UnitPrice * il.Quantity) as TotalSpent\nFROM Customer c\nJOIN Invoice i ON c.CustomerId = i.CustomerId\nJOIN InvoiceLine il ON i.InvoiceId = il.InvoiceId\nGROUP BY c.CustomerId\nORDER BY TotalSpent DESC\nLIMIT 5;"
23
+ },
24
+ {
25
+ "id": 6,
26
+ "status": "FAIL",
27
+ "generated": "SELECT T3.Name AS ArtistName, T1.Name AS SongName\nFROM Track AS T1\nJOIN Album AS T2 ON T1.AlbumId = T2.AlbumId\nJOIN Genre AS T3 ON T1.GenreId = T3.GenreId\nWHERE T3.Name = 'Rock'\nGROUP BY T1.TrackId;"
28
+ },
29
+ {
30
+ "id": 7,
31
+ "status": "ERROR: Execution failed on sql 'SELECT T3.Name AS MediaTypeName, COUNT(T1.TrackId) AS TotalTracksSold\nFROM Track AS T1\nJOIN MediaType AS T2 ON T1.MediaTypeId = T2.MediaTypeId\nGROUP BY T2.Name;': no such column: T3.Name"
32
+ },
33
+ {
34
+ "id": 8,
35
+ "status": "PASS",
36
+ "generated": "SELECT T3.Name FROM Playlist AS T1 INNER JOIN PlaylistTrack AS T2 ON T1.PlaylistId = T2.PlaylistId INNER JOIN Track AS T3 ON T2.TrackId = T3.TrackId WHERE T1.Name = 'TV Shows';"
37
+ },
38
+ {
39
+ "id": 9,
40
+ "status": "PASS",
41
+ "generated": "SELECT a.Name, COUNT(t.TrackId) AS TrackCount\nFROM Artist a\nJOIN Album ar ON a.ArtistId = ar.ArtistId\nJOIN Track t ON ar.AlbumId = t.AlbumId\nGROUP BY a.ArtistId\nORDER BY TrackCount DESC\nLIMIT 1;"
42
+ },
43
+ {
44
+ "id": 10,
45
+ "status": "FAIL",
46
+ "generated": "SELECT g.Name FROM Genre g JOIN Track t ON g.GenreId = t.GenreId GROUP BY g.Name HAVING COUNT(t.TrackId) > 100;"
47
+ },
48
+ {
49
+ "id": 11,
50
+ "status": "FAIL",
51
+ "generated": "SELECT a.Name, SUM(il.UnitPrice * il.Quantity) as TotalRevenue\nFROM Artist a\nJOIN Album al ON a.ArtistId = al.ArtistId\nJOIN Track tr ON al.AlbumId = tr.AlbumId\nJOIN InvoiceLine il ON tr.TrackId = il.TrackId\nGROUP BY a.ArtistId\nORDER BY TotalRevenue DESC\nLIMIT 1;"
52
+ },
53
+ {
54
+ "id": 12,
55
+ "status": "FAIL",
56
+ "generated": "SELECT a.Title, AVG(t.Milliseconds) / 1000 AS AverageLengthInSeconds\nFROM Album a\nJOIN Track t ON a.AlbumId = t.AlbumId\nGROUP BY a.AlbumId;"
57
+ },
58
+ {
59
+ "id": 13,
60
+ "status": "PASS",
61
+ "generated": "SELECT c.FirstName, c.LastName, e.FirstName, e.LastName\nFROM Customer c\nJOIN Employee e ON c.SupportRepId = e.EmployeeId\nWHERE e.LastName = 'Peacock';"
62
+ },
63
+ {
64
+ "id": 14,
65
+ "status": "PASS",
66
+ "generated": "SELECT t2.City FROM Invoice AS t1 INNER JOIN Customer AS t2 ON t1.CustomerId = t2.CustomerId WHERE strftime('%Y', t1.InvoiceDate) = '2013' GROUP BY t2.City ORDER BY COUNT(*) DESC LIMIT 1;"
67
+ },
68
+ {
69
+ "id": 15,
70
+ "status": "ERROR: Execution failed on sql 'SELECT T1.Title FROM Album AS T1 INNER JOIN InvoiceLine AS T2 ON T1.AlbumId = T2.AlbumId GROUP BY T1.AlbumId HAVING SUM(T2.UnitPrice * T2.Quantity) > 20;': no such column: T2.AlbumId"
71
+ }
72
+ ]
requirements.txt ADDED
Binary file (492 Bytes). View file
 
src/database/Chinook_Sqlite.sqlite ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7651ba378ac2fcd0dfc3c66fb101f7a7eed3ba39a612ec642b96e20702061f15
3
+ size 1007616
src/database/__pycache__/db_manager.cpython-313.pyc ADDED
Binary file (2.56 kB). View file
 
src/database/db_manager.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # This module provides a function to establish a connection to the SQLite database used in the NL2SQL project. It also includes a test block to verify the connection and list the tables in the database.
2
+
3
+ import sqlite3
4
+ import os
5
+
6
+ # Get the path to the database file
7
+ DB_PATH = os.path.join(os.path.dirname(__file__), 'Chinook_Sqlite.sqlite')
8
+
9
+ def get_db_connection():
10
+ """Establishes a connection to the SQLite database."""
11
+ try:
12
+ connection = sqlite3.connect(DB_PATH)
13
+ return connection
14
+ except sqlite3.Error as e:
15
+ print(f"Error connecting to database: {e}")
16
+ return None
17
+
18
+ # Test the database connection
19
+ if __name__ == "__main__":
20
+ connection = get_db_connection()
21
+ if connection:
22
+ print("Database connection successful!")
23
+ cursor = connection.cursor()
24
+ cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
25
+ print("Tables in the database:", cursor.fetchall())
26
+ connection.close()
27
+ else:
28
+ print("Failed to connect to the database.")
29
+
30
+ # Extract Schema Information for LLM Prompts
31
+ def get_schema_context():
32
+ """Extracts the database schema information to be used in LLM prompts."""
33
+ connection = get_db_connection()
34
+ if not connection:
35
+ return "Unable to connect to the database to retrieve schema information."
36
+
37
+ cursor = connection.cursor()
38
+ cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
39
+ tables = [t[0] for t in cursor.fetchall() if not t[0].startswith('sqlite_')]
40
+
41
+ schema_text = ""
42
+ for table in tables:
43
+ cursor.execute(f"PRAGMA table_info({table});")
44
+ columns = [f"{c[1]} ({c[2]})" for c in cursor.fetchall()]
45
+ schema_text += f"Table {table}: {', '.join(columns)}\n"
46
+ connection.close()
47
+ return schema_text
48
+
src/nl2sql/__pycache__/qwen_engine.cpython-313.pyc ADDED
Binary file (2.4 kB). View file
 
src/nl2sql/qwen_engine.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Inference Script for Qwen-2.5:1B
2
+
3
+ import requests
4
+ import json
5
+
6
+ def query_qwen(prompt):
7
+ url = "http://localhost:11434/api/generate"
8
+ payload = {
9
+ "model": "qwen2.5-coder:1.5b",
10
+ "prompt": prompt,
11
+ "stream": False,
12
+ "options": {
13
+ # Set temperature to 0 for deterministic output
14
+ "temperature": 0
15
+ }
16
+ }
17
+ response = requests.post(url, json=payload)
18
+ raw_response = response.json().get("response", "").strip()
19
+
20
+ # Clean the response to extract only the SQL query
21
+ if raw_response.startswith("```sql") and raw_response.endswith("```"):
22
+ return raw_response.replace("```sql", "").replace("```", "").strip()
23
+ else:
24
+ return raw_response
25
+
26
+ def build_qwen_prompt(question, schema):
27
+ return f"""## Instructions:
28
+ You are an expert SQLite assistant that converts natural language questions into SQL queries.
29
+
30
+ Use ONLY SQLite syntax.
31
+
32
+ Important database relationships:
33
+ - Artist(ArtistId) -> Album(ArtistId)
34
+ - Album(AlbumId) -> Track(AlbumId)
35
+ - Track(GenreId) -> Genre(GenreId)
36
+ - Track(MediaTypeId) -> MediaType(MediaTypeId)
37
+ - Track(TrackId) -> InvoiceLine(TrackId)
38
+ - InvoiceLine(InvoiceId) -> Invoice(InvoiceId)
39
+ - Customer(CustomerId) -> Invoice(CustomerId)
40
+ - Customer(SupportRepId) -> Employee(EmployeeId)
41
+ - Playlist(PlaylistId) -> PlaylistTrack(PlaylistId)
42
+ - PlaylistTrack(TrackId) -> Track(TrackId)
43
+
44
+ SQL Rules:
45
+ 1. Use JOIN when data comes from multiple tables.
46
+ 2. Use GROUP BY when using COUNT, SUM, or AVG.
47
+ 3. Use HAVING for filtering aggregated values.
48
+ 4. Use ORDER BY for ranking results.
49
+ 5. Use LIMIT instead of TOP.
50
+ 6. Always produce executable SQLite SQL.
51
+ 7. Always use table aliases (t, a, ar, g, etc.) when joining tables.
52
+
53
+ Return ONLY the SQL query.
54
+ Do not provide explanations.
55
+
56
+ ## Database Schema:
57
+ {schema}
58
+
59
+ ## Question:
60
+ {question}
61
+
62
+ ## SQL Query:
63
+ """
src/scripts/__pycache__/evaluate.cpython-313.pyc ADDED
Binary file (3.31 kB). View file
 
src/scripts/evaluate.py ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Evaluation Script for NL2SQL Models
2
+
3
+ import json
4
+ import pandas as pd
5
+ from src.database.db_manager import get_db_connection, get_schema_context
6
+ from src.nl2sql.qwen_engine import query_qwen, build_qwen_prompt
7
+
8
+ def compare_results(df_generated_sql, df_gold_sql):
9
+ """Compares the generated SQL query with the expected SQL query."""
10
+ if df_generated_sql is None or df_gold_sql is None:
11
+ return False
12
+ try:
13
+ return df_generated_sql.values.tolist() == df_gold_sql.values.tolist()
14
+ except Exception as e:
15
+ print(f"Error comparing results: {e}")
16
+ return False
17
+
18
+ def run_evaluation():
19
+ # Load test cases from a JSON file
20
+ with open('src/scripts/test_cases.json', 'r') as f:
21
+ test_cases = json.load(f)
22
+
23
+ results = []
24
+ correct_count = 0
25
+ schema = get_schema_context()
26
+
27
+ print(f"Running evaluation on {len(test_cases)} test cases...\n")
28
+
29
+ for case in test_cases:
30
+ print(f"Testing ID {case['id']}: {case['question'][:50]}...")
31
+
32
+ # Generate SQL query using the LLM
33
+ prompt = build_qwen_prompt(case['question'], schema)
34
+ generated_sql = query_qwen(prompt)
35
+
36
+ # Execute both generated SQL and gold SQL
37
+ connection = get_db_connection()
38
+ try:
39
+ df_generated = pd.read_sql_query(generated_sql, connection)
40
+ df_gold = pd.read_sql_query(case['gold_sql'], connection)
41
+
42
+ # Compare results
43
+ is_correct = compare_results(df_generated, df_gold)
44
+ if is_correct:
45
+ correct_count += 1
46
+
47
+ results.append({
48
+ "id": case['id'],
49
+ "status": "PASS" if is_correct else "FAIL",
50
+ "generated": generated_sql
51
+ })
52
+ except Exception as e:
53
+ results.append({"id": case['id'], "status": f"ERROR: {str(e)}"})
54
+ finally:
55
+ connection.close()
56
+
57
+ # Print summary
58
+ accuracy = (correct_count / len(test_cases)) * 100
59
+ print(f"\n EVALUATION COMPLETE")
60
+ print(f"Total Test Cases: {len(test_cases)}")
61
+ print(f"Correctly Generated SQL: {correct_count}")
62
+ print(f"Execution Accuracy: {accuracy:.2f}%")
63
+
64
+ # Save results to a JSON file
65
+ with open('evaluation_results.json', 'w') as f:
66
+ json.dump(results, f, indent=4)
67
+
68
+ if __name__ == "__main__":
69
+ run_evaluation()
src/scripts/test_cases.json ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [
2
+ {
3
+ "id": 1,
4
+ "question": "How many tracks are there in each genre? List the genre name and the count.",
5
+ "gold_sql": "SELECT t.Genre, COUNT(t.TrackId) AS TrackCount FROM Track t GROUP BY t.Genre;"
6
+ },
7
+ {
8
+ "id": 2,
9
+ "question": "Provide a list of all albums and the name of the artist who created them.",
10
+ "gold_sql": "SELECT a.Title, ar.Name FROM Album a JOIN Artist ar ON a.ArtistId = ar.ArtistId;"
11
+ },
12
+ {
13
+ "id": 3,
14
+ "question": "What is the total revenue generated from each country?",
15
+ "gold_sql": "SELECT BillingCountry, SUM(Total) FROM Invoice GROUP BY BillingCountry;"
16
+ },
17
+ {
18
+ "id": 4,
19
+ "question": "Show the full names of all employees who are Sales Support Agents.",
20
+ "gold_sql": "SELECT e.FirstName, e.LastName FROM Employee e WHERE e.Title = 'Sales Support Agent';"
21
+ },
22
+ {
23
+ "id": 5,
24
+ "question": "List the top 5 customers who have spent the most money.",
25
+ "gold_sql": "SELECT c.CustomerId, c.FirstName, c.LastName, SUM(i.Total) as TotalSpent FROM Customer c JOIN Invoice i ON c.CustomerId = i.CustomerId GROUP BY c.CustomerId, c.FirstName, c.LastName ORDER BY TotalSpent DESC LIMIT 5;"
26
+ },
27
+ {
28
+ "id": 6,
29
+ "question": "List all Rock songs and the artists who performed them.",
30
+ "gold_sql": "SELECT t.Name, ar.Name FROM Track t JOIN Genre g ON t.GenreId = g.GenreId JOIN Album a ON t.AlbumId = a.AlbumId JOIN Artist ar ON a.ArtistId = ar.ArtistId WHERE g.Name = 'Rock';"
31
+ },
32
+ {
33
+ "id": 7,
34
+ "question": "Find the total number of tracks sold for each media type.",
35
+ "gold_sql": "SELECT m.Name, COUNT(il.TrackId) FROM MediaType m JOIN Track t ON m.MediaTypeId = t.MediaTypeId JOIN InvoiceLine il ON t.TrackId = il.TrackId GROUP BY m.Name;"
36
+ },
37
+ {
38
+ "id": 8,
39
+ "question": "Show the names of all tracks that appear on the 'TV Shows' playlist.",
40
+ "gold_sql": "SELECT t.Name FROM Track t JOIN PlaylistTrack pt ON t.TrackId = pt.TrackId JOIN Playlist p ON pt.PlaylistId = p.PlaylistId WHERE p.Name = 'TV Shows';"
41
+ },
42
+ {
43
+ "id": 9,
44
+ "question": "Which artist has the most tracks? Give the name and count.",
45
+ "gold_sql": "SELECT ar.Name, COUNT(t.TrackId) FROM Artist ar JOIN Album a ON ar.ArtistId = a.ArtistId JOIN Track t ON a.AlbumId = t.AlbumId GROUP BY ar.Name ORDER BY COUNT(t.TrackId) DESC LIMIT 1;"
46
+ },
47
+ {
48
+ "id": 10,
49
+ "question": "Which genres have more than 100 tracks?",
50
+ "gold_sql": "SELECT g.Name, COUNT(t.TrackId) as TrackCount FROM Genre g JOIN Track t ON g.GenreId = t.GenreId GROUP BY g.Name HAVING TrackCount > 100;"
51
+ },
52
+ {
53
+ "id": 11,
54
+ "question": "Who is the best-selling artist by total revenue? Provide the artist's name and total revenue.",
55
+ "gold_sql": "SELECT ar.Name, SUM(i.Total) as TotalRevenue FROM Artist ar JOIN Album a ON ar.ArtistId = a.ArtistId JOIN Track t ON a.AlbumId = t.AlbumId JOIN InvoiceLine il ON t.TrackId = il.TrackId JOIN Invoice i ON il.InvoiceId = i.InvoiceId GROUP BY ar.Name ORDER BY TotalRevenue DESC LIMIT 1;"
56
+ },
57
+ {
58
+ "id": 12,
59
+ "question": "Find the average length of tracks in seconds for each album. List the album title and average length.",
60
+ "gold_sql": "SELECT a.Title, AVG(t.Milliseconds) as AverageLength FROM Album a JOIN Track t ON a.AlbumId = t.AlbumId GROUP BY a.Title;"
61
+ },
62
+ {
63
+ "id": 13,
64
+ "question": "List customers helped by the employee Jane Peacock. Provide the customer's full name and the employee's full name.",
65
+ "gold_sql": "SELECT c.FirstName AS CustomerFirstName, c.LastName AS CustomerLastName, e.FirstName AS EmployeeFirstName, e.LastName AS EmployeeLastName FROM Customer c JOIN Employee e ON c.SupportRepId = e.EmployeeId WHERE e.FirstName = 'Jane' AND e.LastName = 'Peacock';"
66
+ },
67
+ {
68
+ "id": 14,
69
+ "question": "Which city had the highest number of invoices in 2013?",
70
+ "gold_sql": "SELECT BillingCity, COUNT(InvoiceId) FROM Invoice WHERE InvoiceDate LIKE '2013%' GROUP BY BillingCity ORDER BY 2 DESC LIMIT 1;"
71
+ },
72
+ {
73
+ "id": 15,
74
+ "question": "List albums with a total price greater than 20 dollars.",
75
+ "gold_sql": "SELECT al.Title, SUM(t.UnitPrice) FROM Album al JOIN Track t ON al.AlbumId = t.AlbumId GROUP BY al.Title HAVING SUM(t.UnitPrice) > 20;"
76
+ }
77
+ ]
test_bench.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Test the database connection and schema extraction
2
+
3
+ import pandas as pd
4
+ from src.database.db_manager import get_db_connection, get_schema_context
5
+ from src.nl2sql.qwen_engine import query_qwen, build_qwen_prompt
6
+
7
+ def run_test(question):
8
+ print(f"\n[USER QUESTION]: {question}")
9
+
10
+ # Step 1: Get the database schema context from db_manager
11
+ schema = get_schema_context()
12
+
13
+ # Step 2: Build the prompt and Generate SQL query
14
+ prompt = build_qwen_prompt(question, schema)
15
+ generated_sql = query_qwen(prompt)
16
+ print(f"\n[GENERATED SQL]:\n{generated_sql}")
17
+
18
+ # Step 3: Execute the generated SQL query against the database
19
+ try:
20
+ connection = get_db_connection()
21
+ # Use pandas to execute the SQL query and display results
22
+ result_df = pd.read_sql_query(generated_sql, connection)
23
+ connection.close()
24
+
25
+ print("[EXECUTION SUCCESS]:")
26
+ if result_df.empty:
27
+ print("Query returned no results.")
28
+ else:
29
+ print(result_df.head(5)) # Display the first 5 rows of the result
30
+ except Exception as e:
31
+ print(f"[EXECUTION ERROR]: {e}")
32
+
33
+ if __name__ == "__main__":
34
+ # Example test question
35
+ test_question = "What is the total revenue generated from each country?"
36
+ run_test(test_question)