Spider NL2SQL

An execution-grounded, multi-agent text-to-SQL system for the Spider benchmark. It combines hybrid schema retrieval, foreign-key graph expansion, a SQL-focused language model, structural checks, read-only SQLite execution, and automatic repair.

The original research notebooks remain in this repository. The reusable implementation now lives in an installable Python package with typed configuration, a command-line interface, and tests.

How it works

flowchart LR
    Q[Question + database ID] --> R{Retrieval router}
    R -->|lexical match| B[BM25 schema retrieval]
    R -->|semantic match| D[E5 + FAISS retrieval]
    R -->|low confidence| F[Full schema fallback]
    B --> K[Foreign-key path expansion]
    D --> K
    K --> G[SQL generator]
    F --> G
    G --> S[Shape audit and optional reprompt]
    S --> E[Read-only SQLite execution]
    E -->|error| X[Deterministic / LLM repair]
    X --> E
    E --> O[SQL + JSONL trace]

Loading

The pipeline is designed around one invariant: generated SQL is never treated as successful until SQLite can execute it against the target database.

Features

• Hybrid per-database retrieval with BM25 and E5 embeddings backed by FAISS
• Confidence-based routing with a full-schema fallback
• Foreign-key shortest-path expansion to include bridge tables and join edges
• 4-bit loading for XGenerationLab/XiYanSQL-QwenCoder-7B-2502
• English and Arabic heuristics for aggregation, grouping, uniqueness, and ranking
• Read-only SQLite execution
• Deterministic typo repair followed by an optional LLM repair loop
• Batch predictions plus structured JSONL metadata and traces
• Lightweight unit tests that do not download model weights

Repository layout

.
├── src/nl2sql/                 # Installable application package
│   ├── agents.py               # Generation, shape checks, execution, repair
│   ├── cli.py                  # `nl2sql` command
│   ├── config.py               # Typed configuration
│   ├── data.py                 # Spider loaders
│   ├── pipeline.py             # Composition and orchestration
│   ├── retrieval.py            # BM25, dense retrieval, and routing
│   └── schema.py               # Schema docs and foreign-key graph
├── tests/                      # Fast unit tests
├── final-nl2sql-multi-agent-system.ipynb
├── Helpers/                    # Research and ablation notebooks
├── Multi-agent-NL2SQL/         # Earlier multi-agent experiments
├── Dataset/                    # Archived dataset artifacts
├── pyproject.toml
└── .env.example

Requirements

• Python 3.10 or newer
• A Spider 1.0 checkout containing tables.json and database/
• An NVIDIA GPU is strongly recommended for the default 7B generator
• Enough disk space for the generator and embedding model weights

The repository contains compressed dataset artifacts, not a ready-to-run Spider directory. Extract or download Spider separately and pass its root directory to the CLI.

Installation

Create an isolated environment and install the project:

python -m venv .venv
source .venv/bin/activate            # Windows: .venv\Scripts\activate
python -m pip install --upgrade pip
python -m pip install -e ".[models]"

For development tools:

python -m pip install -e ".[models,dev]"

bitsandbytes is installed only on Linux. On another platform, pass --no-4bit and ensure the model fits in available memory, or run the project in a Linux GPU environment such as Kaggle.

Dataset layout

The path supplied as --spider-root must look like this:

spider/
├── tables.json
├── dev.json
└── database/
    └── concert_singer/
        └── concert_singer.sqlite

An embedding model may be either a Hugging Face model ID such as intfloat/e5-small-v2 or a local Sentence Transformers checkpoint produced by the tuning notebook.

Usage

Batch prediction

nl2sql \
  --spider-root /data/spider \
  --embed-model intfloat/e5-small-v2 \
  predict \
  --input /data/spider/dev.json \
  --output outputs/dev_predictions.sql \
  --metadata outputs/dev_metadata.jsonl

Use --limit 10 for a smoke run. The SQL output keeps exactly one line per input example, including a safe SELECT 1 placeholder if an individual example raises an unexpected exception.

Single question

nl2sql \
  --spider-root /data/spider \
  query \
  --db-id concert_singer \
  --question "How many singers are there?"

Useful global options include --model, --embed-model, --embed-device, --no-4bit, and --no-llm-repair. Place global options before the predict or query subcommand.

Python API

from pathlib import Path

from nl2sql import RuntimeConfig, build_pipeline
from nl2sql.data import load_spider

runtime = RuntimeConfig(
    spider_root=Path("/data/spider"),
    embed_model="intfloat/e5-small-v2",
)
examples, schemas, database_ids = load_spider(
    runtime.spider_root / "dev.json",
    runtime.tables_path,
)
pipeline = build_pipeline(runtime, schemas, database_ids)
result = pipeline.run("How many singers are there?", "concert_singer")

print(result["final_sql"])
print(result["status"], result["router_mode"])

Constructing the pipeline loads model weights and builds indexes for the supplied database IDs. For interactive workloads, construct it once and reuse it.

Outputs

The batch command creates:

• *.sql: one predicted SQL statement per line, compatible with Spider evaluation.
• *.jsonl: one record per example containing status, repair method, row count, router mode, remaining shape issues, runtime error, and a stage-by-stage trace.

SQLite connections use read-only mode. This prevents generated mutation statements from changing benchmark databases.

Evaluation

Use the official Spider evaluator from its repository:

git clone https://github.com/taoyds/spider.git vendor/spider
python vendor/spider/evaluation.py \
  --gold /data/spider/dev_gold.sql \
  --pred outputs/dev_predictions.sql \
  --db /data/spider/database \
  --table /data/spider/tables.json \
  --etype all

Spider reports structural exact/partial matching and execution accuracy. Inspect the JSONL trace alongside aggregate metrics to distinguish retrieval, generation, shape, and execution failures.

Development

Run the fast test suite and linter before committing:

pytest
ruff check src tests

The tests exercise SQL cleanup, shape detection, deterministic repair, schema documents, and foreign-key expansion without requiring a GPU or network access.

Research notebooks

The notebooks are retained as reproducible research artifacts:

Location	Purpose
final-nl2sql-multi-agent-system.ipynb	Canonical end-to-end experiment used to derive the package
Helpers/Ranker_Fine_Tuning.ipynb	Spider-specific E5 retrieval tuning
Helpers/value-linker.ipynb	Value-linking experiments
Helpers/hybrid-router-v2.ipynb	Hybrid routing experiments
Helpers/agents_without_tuning.ipynb	Untuned embedding baseline
Helpers/nl2sql-rag-self-repair-hr-value-linker-agent.ipynb	RAG, repair, and high-recall value-linking experiment
Multi-agent-NL2SQL/	Structured-message and function-calling variants

Notebook paths still target Kaggle and may contain exploratory duplication. Use the package for repeatable runs and the notebooks for experiment history.

Additional project background is available in NL2SQL_project_documentation.pdf and Natural Language to SQL (NL2SQL).pptx.

Current limitations

• Indexes are rebuilt in memory on each process start.
• Shape checks are heuristic and can produce false positives.
• Successful execution does not guarantee semantic correctness.
• The default generator is resource-intensive and has not been abstracted behind a remote inference provider.
• The official Spider evaluator is intentionally not vendored.