Views
No views yet
nnul/sqlchat: A Conversational AI for SQL Generationsqlchat, a powerful and efficient language model designed specifically for Text-to-SQL tasks. It can understand natural language questions and database schemas to generate accurate SQL queries, including complex statements for creating and managing tables (Data Definition Language).CREATE TABLE contexts provided in the prompt to generate relevant queries.CREATE TABLE statements, including constraints like PRIMARY KEY and FOREIGN KEY relationships.JOINs, aggregations (COUNT, MAX), and sorting (ORDER BY ... LIMIT).sqlchat is with the Unsloth library, which will ensure you get the best performance.1pip install unsloth
2pip install "torch>=2.3.1"1import torch
2from unsloth import FastLanguageModel
3from transformers import TextStreamer
4
5# Load the sqlchat model from the Hugging Face Hub
6# This is a standalone 4-bit model, so we load it as such.
7print("Loading sqlchat model...")
8model, tokenizer = FastLanguageModel.from_pretrained(
9 model_name="nnul/sqlchat",
10 max_seq_length=4096,
11 dtype=None,
12 load_in_4bit=True,
13)
14print("Model loaded successfully.")
15
16# This call optimizes the model for the fastest possible inference.
17FastLanguageModel.for_inference(model)
18
19def generate_sql(instruction: str, context: str = ""):
20 """
21 A helper function to generate SQL from a natural language prompt.
22 """
23 prompt = tokenizer.apply_chat_template(
24 [
25 {"role": "system", "content": "You are a helpful assistant that generates SQL queries based on natural language questions and database schemas."},
26 {"role": "user", "content": f"### Instruction:\n{instruction}\n\n### Context:\n{context}"},
27 ],
28 tokenize=False,
29 add_generation_prompt=True,
30 enable_thinking=False, # Ensures direct SQL output
31 )
32
33 inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
34 text_streamer = TextStreamer(tokenizer, skip_prompt=True, clean_up_tokenization_spaces=True)
35
36 print(f"User Instruction: {instruction}")
37 print("\nModel Output:")
38 print("---------------------------------")
39 _ = model.generate(
40 **inputs,
41 streamer=text_streamer,
42 max_new_tokens=256,
43 do_sample=False, # Use greedy decoding for deterministic output
44 use_cache=True,
45 )
46 print("---------------------------------\n")
47
48# --- Example 1: Querying Data ---
49generate_sql(
50 instruction="Which department has the most number of employees?",
51 context="CREATE TABLE department (name VARCHAR, num_employees INTEGER)"
52)
53
54# --- Example 2: Creating a Table (DDL) ---
55generate_sql(
56 instruction="We need a table to manage student enrollments in courses. This table should link the 'students' table and the 'courses' table using their respective IDs.",
57 context="""
58CREATE TABLE students (student_id INTEGER PRIMARY KEY, student_name VARCHAR(255));
59CREATE TABLE courses (course_id INTEGER PRIMARY KEY, course_title VARCHAR(255));
60"""
61)User Instruction: Which department has the most number of employees?
Model Output:
---------------------------------
SELECT name FROM department ORDER BY num_employees DESC LIMIT 1;
---------------------------------
User Instruction: We need a table to manage student enrollments in courses. This table should link the 'students' table and the 'courses' table using their respective IDs.
Model Output:
---------------------------------
CREATE TABLE student_enrollment (student_id INTEGER, course_id INTEGER, PRIMARY KEY (student_id, course_id), FOREIGN KEY (student_id) REFERENCES students(student_id), FOREIGN KEY (course_id) REFERENCES courses(course_id));
---------------------------------<|im_start|>system
You are a helpful assistant that generates SQL queries based on natural language questions and database schemas.<|im_end|>
<|im_start|>user
### Instruction:
{Your natural language question}
### Context:
{The CREATE TABLE statements for the relevant tables}<|im_end|>
<|im_start|>assistant