keshavsharma/tinyllama-1.1b-sql-reasoning

TinyLlama-1.1B Text-to-SQL — LoRA SFT with Reasoning

A parameter-efficient fine-tuned adapter for TinyLlama/TinyLlama-1.1B-Chat-v1.0 that converts natural-language questions and SQL database schemas into SQL queries, trained to produce an explicit reasoning trace before the final answer.

This is a second-generation adapter in the same series as keshavsharma/qwen2.5-3b-sql-lora-sft and an earlier plain-SQL TinyLlama adapter. Two changes from the earlier TinyLlama attempt drove a large jump in output quality: switching the training set from b-mc2/sql-create-context to gretelai/synthetic_text_to_sql (broader task-type and complexity coverage, including joins and analytics-style questions), and training the model to reason step-by-step inside <reasoning> tags before committing to a final <answer>.

Model Details

Intended Use

The model takes:

• —A database schema, as CREATE TABLE statements (SQL Context).
• —A natural-language question about the database.

It generates a short reasoning trace identifying the relevant tables/columns and query shape, followed by the final SQL query.

Output Format

Unlike a plain text-to-SQL model, this adapter is trained to always respond with a reasoning block followed by an answer block:

<reasoning>...step-by-step derivation of the query...</reasoning><answer>...final SQL query...</answer>

Downstream code must parse the `<answer>...</answer>` span out of the generation before treating it as executable SQL — do not execute the raw model output directly.

Example

Schema:

CREATE TABLE employees (
    employee_id INTEGER,
    name TEXT,
    salary REAL,
    department TEXT
);

Question: What is the average salary of employees in the engineering department?

Generated output:

<reasoning>This query calculates the average salary of employees in the engineering department by using the AVG function on the salary column, and filtering the employees table for rows where the department is 'Engineering'.</reasoning><answer>SELECT AVG(salary) FROM employees WHERE department = 'Engineering';</answer>

Training Method

Supervised Fine-Tuning (SFT) with LoRA. Instead of updating all parameters of the base model, LoRA adds trainable low-rank matrices to selected transformer layers while keeping the original weights frozen.

Loss Masking

Loss is computed only on the reasoning + answer tokens (<reasoning>...</reasoning><answer>...</answer>). System prompt, schema, and question tokens are masked with -100 in the labels, so the model is trained to produce the reasoning-and-answer continuation rather than being penalized for reproducing the input prompt.

Prompt Format

Uses TinyLlama's chat template. The system instruction used during training:

You are a text-to-SQL assistant. Given a database schema (CREATE TABLE statements) and a question, write the SQL query that answers the question. Respond with only the SQL query. Reason with <reasoning></reasoning> blocks. Answer with <answer></answer> blocks.

The user turn follows the format:

SQL Context: {schema}

Question: {question}

Write the SQL query that answers the question.

Installation

pip install torch transformers datasets accelerate peft

Inference

import re
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
ADAPTER = "keshavsharma/tinyllama-1.1b-sql-reasoning"

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL, torch_dtype=torch.bfloat16, device_map="auto"
)
model = PeftModel.from_pretrained(base_model, ADAPTER)
model.eval()
model.config.use_cache = True

SYSTEM_PROMPT = (
    "You are a text-to-SQL assistant. Given a database schema (CREATE TABLE "
    "statements) and a question, write the SQL query that answers the question. "
    "Respond with only the SQL query. Reason with <reasoning></reasoning> blocks. "
    "Answer with <answer></answer> blocks."
)

def generate_sql(schema, question):
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"SQL Context: {schema}\n\nQuestion: {question}\n\nWrite the SQL query that answers the question."},
    ]
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs, max_new_tokens=768, do_sample=False, pad_token_id=tokenizer.eos_token_id
        )
    raw = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()

    match = re.search(r"<answer>(.*?)</answer>", raw, re.DOTALL)
    sql = match.group(1).strip() if match else None
    return {"raw": raw, "sql": sql}

Note the reference implementation used during training generated with do_sample=False for evaluation, but the training-time sanity-check cell in this repo's notebook used sampling (top_p=0.95, temperature=0.9) with do_sample=False set simultaneously — those sampling parameters are ignored when do_sample=False, so generation was effectively greedy/deterministic throughout. Keep do_sample=False for reproducible evaluation; only enable sampling deliberately if you want varied reasoning traces.

Limitations

• —Generated SQL should be validated — and ideally checked that every referenced table/column actually exists in the given schema — before execution. Spot testing showed occasional hallucinated column names on harder correlated-subquery questions (e.g. inventing a column that doesn't appear in the schema).
• —Correlated aggregate arithmetic inside HAVING (e.g. comparing an average against a per-row-derived threshold computed from another aggregate) remains a weak spot; treat multi-step nested-aggregate questions with extra caution.
• —As a 1.1B-parameter base model, this adapter has materially less headroom than larger Text-to-SQL fine-tunes (e.g. the Qwen2.5-3B version in the same series) on harder multi-table, multi-condition questions, even with reasoning traces.
• —The reasoning trace reflects the model's stated derivation, not a verified execution trace — it can be fluent and plausible while the final SQL is still wrong. Don't treat the presence of correct-sounding reasoning as confirmation the SQL is correct.
• —No claim of production-level Text-to-SQL reliability is made. Outputs should be reviewed or executed against a real schema before use in any automated pipeline.

Citation / References

• —TinyLlama — TinyLlama: An Open-Source Small Language Model
• —LoRA — Hu et al., LoRA: Low-Rank Adaptation of Large Language Models
• —PEFT — Hugging Face Parameter-Efficient Fine-Tuning
• —Synthetic Text-to-SQL dataset — gretelai/synthetictextto_sql