CHASE-SQL

Reproduction of CHASE-SQL: Multi-Path Reasoning and Preference-Optimized Candidate Selection in Text-to-SQL (ICLR 2025, arXiv:2410.01943).

Evaluated on the BIRD benchmark dev set (1,534 questions).

---

Pipeline (4 Stages)

Question + Evidence + DB
        │
        ▼
  1. Value Retrieval      ──  LSH + reranking to find exact DB values
        │
        ▼
  2. Candidate Generation ──  DC-CoT + QP-CoT + OS-ICL  (3 × 7 = 21 candidates)
        │
        ▼
  3. Query Fixer          ──  β=3 self-repair loop, drops unfixable candidates
        │
        ▼
  4. Selection Agent      ──  Pairwise tournament (both orderings per pair)
        │
        ▼
     Final SQL

---

Repository Structure

├── pipeline.py                  # Main entry point — runs all 4 stages
├── src/
│   ├── value_retrieval.py       # Stage 1 — keyword extraction + LSH + reranking
│   ├── candidate_gen.py         # Stage 2 — DC-CoT, QP-CoT, OS-ICL generators
│   ├── query_fixer.py           # Stage 3 — self-repair loop (Algorithm 4)
│   ├── selection.py             # Stage 4 — pairwise tournament (Algorithm 3)
│   ├── llm.py                   # LLM wrapper (Gemini / Claude / Mistral)
│   ├── selector_llm.py          # Selector wrapper (5 modes)
│   ├── db_utils.py              # SQLite execution + schema serialization
│   └── logger.py                # Structured logging
├── scripts/
│   ├── preprocess_lsh.py        # Build LSH indexes (run once before inference)
│   └── evaluate.py              # Runs the full pipeline and measures accuracy
├── fine_tuning/
│   ├── generate_pairs.py        # Generate selector training pairs from BIRD train
│   ├── train_selector_gemma.py  # Gemma 2 9B — LoRA on Wulver HPC
│   ├── train_selector_qwen.py   # Qwen2.5-Coder — LoRA on Wulver HPC
│   ├── train_selector_flash.py  # Gemini 1.5 Flash — SFT on Vertex AI
│   ├── evaluate_selector.py     # Binary accuracy eval for all 5 selector modes
│   └── slurm_finetune.sh        # SLURM job script for Wulver
├── configs/
│   ├── config.yaml              # Runtime config (models, paths, hyperparams)
│   └── finetune_config.yaml     # Fine-tuning hyperparams (Gemma / Qwen / Flash)
├── .env.example                 # API key template
└── requirements.txt

---

Setup

Step 1 — Install dependencies

pip install -r requirements.txt

Step 2 — Set API keys

cp .env.example .env

Edit .env with your keys:

GEMINI_API_KEY=...          # Required for Config 1 (generation + selector)
ANTHROPIC_API_KEY=...       # Required for Config 2 (generation)
MISTRAL_API_KEY=...         # Required for Config 3 (generation)
GOOGLE_CLOUD_PROJECT=...    # Required for Flash fine-tuning on Vertex AI

Step 3 — Download BIRD dataset

Download train + dev from bird-bench.github.io and place under data/bird/:

data/bird/
├── train/
│   ├── train.json
│   ├── tables.json
│   └── train_databases/{db_id}/{db_id}.sqlite
└── dev/
    ├── dev.json
    ├── tables.json
    └── dev_databases/{db_id}/{db_id}.sqlite

Step 4 — Build LSH indexes (one-time preprocessing)

python scripts/preprocess_lsh.py

Builds MinHash LSH indexes for every database (train + dev) into data/preprocessed/. ~10-15 min. Use --split dev to build only the dev split.

---

Running the Pipeline

After LSH preprocessing, evaluate.py handles everything — it loads BIRD questions and runs all 4 stages (Value Retrieval → Candidate Generation → Query Fixer → Selection) automatically for each question.

Choose the model

Edit configs/config.yaml before running:

# Which LLM generates SQL candidates (Stages 1-3)
generation_model: gemini-1.5-pro       # or: claude-3.5-sonnet, mistral-large

# Which selector picks the best candidate (Stage 4)
selector_mode: zeroshot_gemini          # or: zeroshot_claude

First time? Use zeroshot_gemini or zeroshot_claude as the selector — these work out of the box with just an API key. Fine-tuned selectors (finetuned_qwen, finetuned_gemma, finetuned_flash) require training first — see the Fine-Tuning the Selector section below.

Run

# Quick smoke test — 100 questions
python scripts/evaluate.py --limit 100

# Full dev set (1,534 questions)
python scripts/evaluate.py

# Override selector from command line (no need to edit config)
python scripts/evaluate.py --selector_mode zeroshot_claude

Results saved to results/dev_predictions_{selector_mode}.json with per-difficulty breakdown (simple / moderate / challenging).

NJIT HPC workflow for gpt-oss-120b + finetuned_qwen

For this combo, do not run --phase all. It can load the generation model and Qwen selector into the same job and trigger GPU OOM.

Use the pipeline in two steps:

# 1. Start vLLM for generation (stages 1-3)
export VLLM_BASE_URL=http://localhost:8000/v1
python scripts/evaluate.py --split dev --limit 100 --phase generate

# 2. Stop vLLM and free the GPUs used by gpt-oss-120b

# 3. Run selection only with the fine-tuned Qwen selector
python scripts/evaluate.py --split dev --limit 100 --phase select --selector_mode finetuned_qwen

This repository now defaults local generation to vllm via local_backend: vllm in configs/config.yaml.

---

Three Configurations (Paper Table 3)

Config	Generation Model	Selector	Dev EX Target
1	gemini-1.5-pro	fine-tuned Flash	>73.01%
2	claude-3.5-sonnet	reuse Config 1 Flash	>69.53%
3	mistral-large	fine-tuned Qwen	>70.33%

Switch by editing configs/config.yaml:

# Config 1
generation_model: gemini-1.5-pro
selector_mode: finetuned_flash

# Config 2
generation_model: claude-3.5-sonnet
selector_mode: finetuned_flash      # reuse — no retraining

# Config 3
generation_model: mistral-large
selector_mode: finetuned_qwen

---

Fine-Tuning the Selector

1. Generate training pairs (~3,800 pairs from BIRD train)

python fine_tuning/generate_pairs.py --config configs/config.yaml

Output: data/finetune/train_pairs.jsonl + val_pairs.jsonl (80/20 split).

2. Train selector models

Gemma 2 9B (LoRA, 4-bit, on Wulver HPC):

python fine_tuning/train_selector_gemma.py --config configs/finetune_config.yaml
# Or via SLURM:
MODEL=gemma sbatch fine_tuning/slurm_finetune.sh

Qwen2.5-Coder (LoRA, 4-bit, on Wulver HPC):

python fine_tuning/train_selector_qwen.py --config configs/finetune_config.yaml
# Or via SLURM:
MODEL=qwen sbatch fine_tuning/slurm_finetune.sh

Gemini 1.5 Flash (Vertex AI SFT):

python fine_tuning/train_selector_flash.py --config configs/finetune_config.yaml
# Copy output tuned_model_name → configs/config.yaml: vertex_tuned_model_name

3. Evaluate all 5 selectors

python fine_tuning/evaluate_selector.py

Output: results/selector_comparison.json

Selector	Binary Accuracy (Paper Table 5)
zeroshot_claude	~60.21%
zeroshot_gemini	~63.98%
finetuned_gemma	>64.28%
finetuned_flash	>71.01%
finetuned_qwen	>58%

---

End-to-End Workflow (all commands in order)

# 1. Setup
pip install -r requirements.txt
cp .env.example .env                          # fill in API keys

# 2. Build LSH indexes (one time)
python scripts/preprocess_lsh.py

# 3. Run pipeline on dev set with zero-shot selector (no fine-tuning needed)
python scripts/evaluate.py --limit 100        # quick test
python scripts/evaluate.py                    # full dev set (1,534 questions)

# 4. Fine-tune selector (optional — for higher accuracy)
python fine_tuning/generate_pairs.py
python fine_tuning/train_selector_qwen.py --config configs/finetune_config.yaml
python fine_tuning/evaluate_selector.py --mode finetuned_qwen

# 5. Full evaluation with fine-tuned selector
python scripts/evaluate.py --selector_mode finetuned_qwen

---

References

• Paper: CHASE-SQL: Multi-Path Reasoning and Preference-Optimized Candidate Selection in Text-to-SQL (ICLR 2025)
• Benchmark: BIRD-SQL