Schema retriever for the enterprise text-to-SQL benchmark

Question in, the 12 tables of the schema ranked out. all-MiniLM-L6-v2 fine-tuned on question/table pairs from the train split, evaluated on the 453 test questions, whose SQL templates never appear in training. Two fine-tuning recipes were tried; the one at the root of this repo is the better of the two (contrastive), the other is under variants/.

Try it in the browser, against the un-tuned model.

Why

In the text-to-SQL project I measured that handing the SQL model a retrieved subset of the schema instead of the full schema made it worse: 50.99 % → 41.72 % execution accuracy for the fine-tuned model. I wrote that up as "retrieval hurt" and moved on, but that conflates two things: retrieval as an idea, and the particular retriever I used. This repo looks at the retriever on its own.

Numbers

"all in top k" is the fraction of test questions for which every table the gold SQL touches is inside the top k. That is the number that matters if you prune the schema before generation, because one missing table means the query cannot be right. recall@k is the softer per-table average.

The labelled contrastive fine-tune beats the off-the-shelf model at every k, and the gap at top 3 (0.963 vs 0.812) is the one that matters for pruning.

The thing I got wrong first

My first attempt was the textbook recipe: MultipleNegativesRankingLoss, (question, table) pairs, other items in the batch as negatives. It improved recall@1 and MRR and made top-3 coverage worse (0.748 vs 0.812 off the shelf).

The reason is the corpus size. With 12 documents, a batch of 12 pairs has every table in it once, so for a question that needs orders and customers, the pair (q, orders) sees customers as a negative. The loss pulls the question toward its first table and pushes it away from its second. Exactly the wrong thing for multi-table questions, which are 46 % of the test set.

The second recipe labels pairs explicitly: every gold table is a positive, five sampled non-gold tables are negatives, OnlineContrastiveLoss, no in-batch negatives. Same base model, same data.

Fine-tuned model (contrastive), all gold tables in the top 3, by slice:

At exactly k = number of tables the question needs (the strictest reading): 1-table questions 0.951, 2-table questions 0.859, 3-table questions 0.852.

Remaining misses at top 3: 17 of 453, listed in misses_top3.json. Sixteen of them are the same failure: "List products held at Marseille" needs inventory and warehouses, but nothing in the question says "warehouse", a city name is doing that job, and the retriever has no way to know Marseille is a warehouse location rather than a customer or supplier one. That is a schema-knowledge problem, not an embedding problem, and the fix is probably to put example values into the table docs.

Training

• base: sentence-transformers/all-MiniLM-L6-v2
• documents: one short text per table, table <name>: columns <col, col, ...> plus the column comments from schema_context.txt; see table_docs.json. No synthetic descriptions.
• MNRL variant: 4,551 pairs, batch 12 with the no-duplicates sampler, 3 epochs, 44.8 min on a laptop CPU
• contrastive variant: 15,216 labelled pairs, batch 32, 2 epochs, margin 0.5, 99.4 min
• lr 3e-5, 10 % warmup, seed 42 for both; evaluation is cosine similarity between the question and the 12 docs, no reranking

Use

from sentence_transformers import SentenceTransformer
import json
from huggingface_hub import hf_hub_download

repo = "hari-krishna-ai/text2sql-schema-retriever-minilm"
m = SentenceTransformer(repo)
docs = json.load(open(hf_hub_download(repo, "table_docs.json")))
tables = sorted(docs)
D = m.encode([docs[t] for t in tables], normalize_embeddings=True)
q = m.encode("Which suppliers ship products that are out of stock?", normalize_embeddings=True)
print(sorted(zip(D @ q, tables), reverse=True)[:5])

ONNX (int8) is in onnx/ for Transformers.js: pipeline('feature-extraction', repo, { dtype: 'q8' }) with pooling: 'mean', normalize: true.

Files

• the shipped sentence-transformers model at the root (modules.json, 1_Pooling/, weights, tokenizer)
• onnx/model_quantized.onnx – int8 export for Transformers.js (the fp32 weights are the safetensors at the root)
• variants/mnrl/ – the other recipe, same layout minus ONNX
• table_docs.json – the 12 documents, exactly as embedded
• metrics.json, misses_top3.json, train.py