slm_text_to_sql/README.md

SLM Text-to-SQL

slm_text_to_sql is a lightweight, local Text-to-SQL translation library powered entirely by a Small Language Model (SLM) running on CPU. It translates natural language questions into database-specific SQL queries.

To achieve production-grade reliability on quantized 1.5B models, the library features context-aware schema pruning and an agentic self-correction loop that validates SQL queries against an in-memory database and corrects execution errors dynamically.

---

Key Features

• Local & Private: Runs completely on CPU / RAM. Zero API keys, zero network latency, and complete data privacy.
• Agentic Self-Correction: Automatically creates an ephemeral in-memory SQLite database, runs the DDL schema, inserts dummy data, and tests the generated query. If execution fails (e.g. syntax, column mismatch), the agent reads the error and dynamically debugs the SQL query in a correction feedback loop.
• Context-Aware Schema Pruning: Prunes large schemas down to a maximum number of relevant tables (default: 8) based on keyword overlap and foreign key relationship paths to prevent model confusion.
• Agnostic & Custom Few-Shots: Overrides/augments few-shot examples with domain-specific templates. Abstract templates are used by default to prevent model copying bias.
• Column Descriptions: Supports passing a glossary dictionary of table/column semantics (e.g. enums, flags) to guide logical mapping.
• Streaming Support: Stream token-by-token output in real-time via a Python generator.

---

Installation

Install locally for development:

# 1. Create a fresh virtual environment
python3 -m venv .venv
source .venv/bin/activate

# 2. Install the package in editable mode
pip install -e .

Note: Requires onnxruntime-genai, huggingface_hub, and pyyaml.

---

Quick Start

from slm_text_to_sql import SLMTextToSQL

# Initialize the agent
agent = SLMTextToSQL()

schema = """
CREATE TABLE users (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    role TEXT DEFAULT 'customer'
);
CREATE TABLE orders (
    id INTEGER PRIMARY KEY,
    customer_id INTEGER REFERENCES users(id),
    total_amount REAL NOT NULL
);
"""

# Generates SQL query and validates/corrects it via SQLite in the background
query = agent.generate_sql(
    schema=schema,
    question="Find the total amount spent by user 'John Doe'."
)

print(query)
# Output:
# SELECT SUM(o.total_amount) FROM orders o JOIN users u ON o.customer_id = u.id WHERE u.name = 'John';

Streaming Example (Bypasses Self-Correction)

from slm_text_to_sql import SLMTextToSQL

agent = SLMTextToSQL()

# Stream tokens as they are generated
for token in agent.generate_sql(
    schema=schema,
    question="List all user names.",
    stream=True
):
    print(token, end="", flush=True)
print()

---

Achieving Great Accuracy (Best Practices)

Small language models (1.5B parameters) are highly sensitive to prompt structure. Follow these design guidelines to achieve 95%+ execution accuracy:

1. Enable Agentic Self-Correction

Keep stream=False (default) to run the self-correction validation loop. If the model hallucinations a column name or a bad join path, the agent captures the exact SQLite error, matches it against table definitions, and feeds it back to the model to correct itself.

2. Document Your Schema DDL (Inline Comments)

Quantized local models rely heavily on semantic hints inside the DDL. Adding inline comments directly to table schemas helps the model map synonyms:

CREATE TABLE users ( -- Represents users, clients, and customers
    id INTEGER PRIMARY KEY,
    status TEXT -- Either 'active' or 'inactive'
);

3. Prevent Few-Shot Name Copying Bias

If you pass custom few-shot examples via few_shot_examples, ensure the table and column names used inside the examples are completely distinct from your actual database tables. Small models tend to prioritize exact strings in examples over DDL schemas if they match.

4. Provide Column Glossaries

Use the column_descriptions parameter to map complex columns or enums to their definitions:

column_descriptions = {
    "users.role": "Can be 'customer', 'manager', or 'admin'",
    "orders.status": "Can be 'pending', 'completed', or 'cancelled'"
}

---

Self-Correction Loop Walkthrough

When a query fails validation, the agent automatically executes the following loop:

graph TD
    A[Initial Generation] --> B[Execute against SQLite DB]
    B -->|Success| C[Return SQL Query]
    B -->|Database Error| D[Extract Table Column Specs]
    D --> E[Construct History & Feedback Prompt]
    E --> F[Correction Attempt]
    F --> B

• Example Database Error: no such column: o.product_id
• Feedback Injected:

  Failed Attempt #1 SQL:
  SELECT p.name FROM products p JOIN orders o ON p.id = o.product_id;
  
  Failed Attempt #1 Database Error:
  no such column: o.product_id
  
  Available column definitions for tables in your query:
  - Table 'products' columns: id, name, price, stock
  - Table 'orders' columns: id, customer_id, total_amount, status
  - Table 'order_items' columns: id, order_id, product_id, quantity
  

• Result: The model immediately notices orders lacks product_id, looks at order_items columns, and joins them correctly.

---

Configuration API

Constructor Configuration

SLMTextToSQL(
    model_path=None,   # Explicit path to an ONNX model directory (optional)
    cache_dir=None,    # Cache directory for auto-downloads
    n_ctx=2048,        # Context window size (default: 2048)
    n_threads=4        # Number of CPU threads (default: 4)
)

Query Generation Method

agent.generate_sql(
    schema: str,                     # Database DDL schema (CREATE TABLE statements)
    question: str,                   # User natural language query
    temperature: float = 0.0,        # Sampling temperature (0.0 for deterministic answers)
    max_tokens: int = None,          # Maximum token limit for the response
    stream: bool = False,            # Returns generator yielding tokens if True
    column_descriptions: dict = None,# Optional dict mapping "table.col" -> "desc"
    few_shot_examples: list = None,  # Custom list of example dicts to append to system prompt
    system_prompt: str = None,       # Overrides the built-in system rules entirely
    max_iterations: int = 5,         # Maximum self-correction retry loops
    max_pruned_tables: int = 8       # Schema pruning table limit cap
)

---

Environment Variables

All constructor parameters can be overridden via environment variables:

Variable	Description	Default
SLM_TEXT_TO_SQL_CONFIG	Path to a custom config.yaml file	—
SLM_TEXT_TO_SQL_CACHE_DIR	Override model download/cache directory	—
SLM_TEXT_TO_SQL_N_THREADS	Number of CPU threads	4
SLM_TEXT_TO_SQL_N_CTX	Context window size	2048
SLM_TEXT_TO_SQL_MAX_TOKENS	Default max tokens per answer	512

---

Fine-Tuning with QLoRA

The library includes a complete script to fine-tune base models (e.g. Qwen/Qwen2.5-Coder-1.5B-Instruct) for Text-to-SQL using QLoRA.

Running Fine-Tuning

python -m slm_text_to_sql.fine_tune

---

License

Apache License 2.0.