Skip to content

Latest commit

 

History

3 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

DB-ER — Database Entity Recognition

Code, result artifacts, and reproduction scripts for Database Entity Recognition (DB-ER): grounding a natural-language question to the tables, columns, and values of a target database, framed as token classification over {O, Table, Column, Value}.

Headline results: a T5-Large tagger reaches 80.8% micro entity F1 in the open-world setting; adding a database-aware MiniLM verifier that grounds predicted spans against the target schema raises this to 93.2% F1.

Publication status

  • Published: Z. Fu, C. Yang, K. Davoudi, K. Q. Pu. "Database Entity Recognition with Data Augmentation and Deep Learning." IEEE IRI 2025, San Jose, CA. DOI 10.1109/IRI66576.2025.00071
  • MSc thesis: Database Entity Recognition using Language Models, Ontario Tech University, 2026.
  • Under review: an extended benchmark paper tobe submitted to the Journal of Web Semantics. It is not accepted or published; please cite the IEEE IRI 2025 paper above.

Data

All annotated splits are on the Hugging Face Hub as Voice49/dber, licensed CC BY 4.0:

Split Rows Role
human_train 500 human-annotated; model/α selection
human_test 500 human-annotated; held-out test
synthetic_train 15,026 generated from paired SQL (AST parsing + ILP)
human_gold_100 100 adjudicated gold subset
Total 16,126

The SQLite databases come from the BIRD and Spider benchmarks and are fetched separately by scripts/download_databases.sh.

Layout

Directory Contents
tagger/ Open-world T5 tagger: training (json_ner_pipeline.py), one-command re-run driver (rerun_train.py), inference + α sweep (inference.py), dataset conversion (convert_dber_to_pipeline.py)
verifier/ Closed-world verifier: db_verification/, CLI (run_db_verification.py)
llm/ LLM comparison (dber_llm/)
analysis/ Bootstrap CIs (bootstrap_pipeline_cis.py), schema extraction for the schema-in-prompt condition (extract_schemas.py)
scripts/ Dataset download (download_dataset.py), database download/layout, verification driver
data/ Committed result artifacts (see data/README.md)
tests/ CPU-only smoke tests for the reproduction path

Setup

python3 -m venv venv && source venv/bin/activate
pip install -r requirements.txt          # training needs a CUDA GPU

Steps 1, 4 and 5 are CPU-only. Step 2 (training) needs a CUDA GPU with ~24 GB of memory for T5-Large; step 3 needs the benchmark databases on disk.

Reproducing the pipeline

1. Data (CPU) — fetch the four splits and convert them to the trainer's format:

python3 scripts/download_dataset.py --out_dir tagger/data
python3 tagger/convert_dber_to_pipeline.py --jsonl_dir tagger/data --out_dir tagger/data

2. Train + predict (GPU):

python3 tagger/rerun_train.py --dev human_train

Trains T5-Large on synthetic_train, trains the linear head, sweeps the recall weight α on human_train, selects α*, and writes boosted prediction JSON for human_test and the gold set. Each step skips itself if its output already exists, so the script is safe to re-run after an interruption.

α* and epoch selection happen on human_train, which is never used to train the tagger. Pass --dev human_test to reproduce the original select-on-test protocol instead.

3. Verify (CPU, needs the databases):

bash scripts/download_databases.sh
bash scripts/run_verification_conditions.sh <prediction.json>

4. Statistics (CPU):

python3 analysis/bootstrap_pipeline_cis.py <verified.json> \
    --config CLF_Tok__TCV__MiniLM

Question-level bootstrap CIs (B=10,000, seed 42, percentile 95% CI over micro entity F1). Add --llm-results-dir data/llm_results with a gold-set input for paired pipeline-vs-LLM deltas.

5. LLM comparison (CPU, needs OPENROUTER_API_KEY):

python3 -m dber_llm.run_experiment --model <slug>
python3 analysis/extract_schemas.py            # for the schema-in-prompt condition
python3 -m dber_llm.run_experiment --model <slug> --schema-file <schemas.json>

Tests

CPU-only, no GPU / database / network required. They check the dataset setup, no missing files, confirms reproduction.

pip install pytest
pytest -q

License

Code in this repository is released under the MIT License (see LICENSE). The DB-ER dataset on the Hugging Face Hub is licensed CC BY 4.0. The BIRD and Spider databases remain under their original licenses.

About

Database Entity Recognition: grounding natural-language questions to database tables, columns, and values. Code, trained-model artifacts, and results for the DB-ER benchmark (IEEE IRI 2025); dataset on the Hugging Face Hub.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages