Text-to-SQL on the ATIS Flight Database

This project explores translating natural language questions about air travel into SQL queries against the ATIS (Air Travel Information System) flight database using T5-small (60M), either fine-tuned from a pretrained checkpoint or trained from scratch on the ATIS training set.

A key contribution is query-specific schema injection: rather than feeding the entire database schema into the model, only the tables relevant to each natural language query are selected, reducing input length and improving accuracy.

---

Repository Structure

.
├── data/
│   ├── train.nl / train.sql        # Training set: NL queries and gold SQL
│   ├── dev.nl   / dev.sql          # Development set
│   ├── test.nl                     # Test set (no gold SQL)
│   ├── flight_database.schema      # Full DB schema
│   └── simplified_schema.txt       # Simplified schema used during training
│
├── results/                        # Model-generated .sql files
├── records/                        # Pickled DB records (.pkl) for evaluation
├── checkpoints/                    # Saved model weights (.pt)
│
├── train_t5.py                     # T5 training loop (fine-tune or from scratch)
├── test_best_model.py              # Run best T5 checkpoint on dev/test
│
├── load_data.py                    # Dataset class & DataLoader construction
├── t5_utils.py                     # T5 model init, checkpoint save/load, optimizer
├── schema_utils.py                 # Query-specific schema extraction heuristics
├── utils.py                        # DB execution, metrics (EM / F1), record I/O
│
├── evaluate.py                     # CLI evaluation script
├── create_ground_truth.py          # Pre-compute dev-set ground-truth DB records
├── error_analysis.py               # Categorise and compare model errors
│
├── requirements.txt
└── .gitignore

---

Setup

Python 3.10 is recommended. Create and activate a virtual environment, then install dependencies:

conda create -n text2sql python=3.10
conda activate text2sql
pip install -r requirements.txt

A CUDA-capable GPU is strongly recommended for training. The code automatically falls back to CPU if no GPU is available.

---

Data

The dataset is based on the ATIS benchmark. Each sample is a natural language question about flights (e.g., "show me all flights from Boston to Denver on Monday") paired with a ground-truth SQL query over the flight database.

The SQLite database (data/flight_database.db) contains ~20 tables including flight, fare, airline, aircraft, airport, city, state, and more.

---

T5-small

Pre-compute ground-truth records (one-time)

python create_ground_truth.py

This saves records/dev_gt_records.pkl, which speeds up per-epoch evaluation.

Training

Fine-tune pretrained T5-small:

python train_t5.py \
  --finetune \
  --experiment_name optimized_schema \
  --max_n_epochs 30 \
  --learning_rate 5e-5 \
  --batch_size 16 \
  --scheduler_type cosine \
  --num_warmup_epochs 2 \
  --patience_epochs 5

Train T5-small from scratch:

python train_t5.py \
  --experiment_name optimized_schema_scratch \
  --max_n_epochs 100 \
  --learning_rate 1e-3 \
  --batch_size 16 \
  --scheduler_type cosine \
  --num_warmup_epochs 5 \
  --patience_epochs 10

Checkpoints are saved under checkpoints/{ft,scr}_experiments/<experiment_name>/.

Testing the best checkpoint

python test_best_model.py \
  --finetune \
  --experiment_name optimized_schema \
  --batch_size 16

Results are written to results/ and records/.

---

Evaluation

python evaluate.py \
  --predicted_sql  results/t5_ft_optimized_schema_dev.sql \
  --predicted_records records/t5_ft_optimized_schema_dev.pkl \
  --development_sql data/dev.sql \
  --development_records records/dev_gt_records.pkl

Three metrics are reported:

• SQL Exact Match — predicted SQL string matches the gold query exactly.
• Record Exact Match — the set of DB records returned is identical to the gold set.
• Record F1 — token-level F1 over the returned DB records (main metric).

---

Error Analysis

python error_analysis.py

Compares fine-tuned vs. from-scratch T5 predictions and categorises errors (wrong table, wrong condition, invalid SQL, etc.). Results are written to error_analysis_results.txt.

---

Experiments & Results

Pre-generated result files in results/: