sqlchat / README.md
nnul's picture
Update README.md
2d23776 verified
# `nnul/sqlchat`: A Conversational AI for SQL Generation
This repository contains `sqlchat`, 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).
This model is provided as a standalone 4-bit quantized model, optimized for easy deployment and high-performance, low-resource inference. It was built using the [Unsloth](https://github.com/unslothai/unsloth) library to ensure maximum speed and memory efficiency.
## Model Capabilities
* **Natural Language to SQL:** Translates complex English questions into executable SQL queries.
* **Schema-Aware:** Understands `CREATE TABLE` contexts provided in the prompt to generate relevant queries.
* **DDL Generation:** Capable of generating `CREATE TABLE` statements, including constraints like `PRIMARY KEY` and `FOREIGN KEY` relationships.
* **Complex Query Logic:** Successfully handles `JOIN`s, aggregations (`COUNT`, `MAX`), and sorting (`ORDER BY ... LIMIT`).
## How to Use
The easiest way to use `sqlchat` is with the Unsloth library, which will ensure you get the best performance.
### Prerequisites
First, install the necessary libraries.
```bash
pip install unsloth
pip install "torch>=2.3.1"
```
### Running Inference
Here is a simple, reusable Python script to run inference with the model.
```python
import torch
from unsloth import FastLanguageModel
from transformers import TextStreamer
# Load the sqlchat model from the Hugging Face Hub
# This is a standalone 4-bit model, so we load it as such.
print("Loading sqlchat model...")
model, tokenizer = FastLanguageModel.from_pretrained(
model_name="nnul/sqlchat",
max_seq_length=4096,
dtype=None,
load_in_4bit=True,
)
print("Model loaded successfully.")
# This call optimizes the model for the fastest possible inference.
FastLanguageModel.for_inference(model)
def generate_sql(instruction: str, context: str = ""):
"""
A helper function to generate SQL from a natural language prompt.
"""
prompt = tokenizer.apply_chat_template(
[
{"role": "system", "content": "You are a helpful assistant that generates SQL queries based on natural language questions and database schemas."},
{"role": "user", "content": f"### Instruction:\n{instruction}\n\n### Context:\n{context}"},
],
tokenize=False,
add_generation_prompt=True,
enable_thinking=False, # Ensures direct SQL output
)
inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
text_streamer = TextStreamer(tokenizer, skip_prompt=True, clean_up_tokenization_spaces=True)
print(f"User Instruction: {instruction}")
print("\nModel Output:")
print("---------------------------------")
_ = model.generate(
**inputs,
streamer=text_streamer,
max_new_tokens=256,
do_sample=False, # Use greedy decoding for deterministic output
use_cache=True,
)
print("---------------------------------\n")
# --- Example 1: Querying Data ---
generate_sql(
instruction="Which department has the most number of employees?",
context="CREATE TABLE department (name VARCHAR, num_employees INTEGER)"
)
# --- Example 2: Creating a Table (DDL) ---
generate_sql(
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.",
context="""
CREATE TABLE students (student_id INTEGER PRIMARY KEY, student_name VARCHAR(255));
CREATE TABLE courses (course_id INTEGER PRIMARY KEY, course_title VARCHAR(255));
"""
)
```
### Expected Output
```
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));
---------------------------------
```
## Performance
The model was benchmarked on an NVIDIA A40 GPU. In a batch-processing scenario, it achieves a throughput of **~55-70 tokens/second**. Single-prompt latency is well within real-time requirements for interactive applications.
* **Peak VRAM Usage (Inference):** ~6.1 GB
## Prompt Template
To get the best results, your prompts should follow this structure:
```
<|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
```