Views
No views yet

1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3tokenizer = AutoTokenizer.from_pretrained("ContextualAI/ctx-bird-reward-250121")
4model = AutoModelForCausalLM.from_pretrained("ContextualAI/ctx-bird-reward-250121")
5
6# Example inputs
7db_schema = """CREATE TABLE customers (
8 id INT PRIMARY KEY,
9 name VARCHAR(100),
10 region VARCHAR(50),
11 revenue FLOAT
12);"""
13
14question = "Show me top 5 highest revenue customers by region"
15evidence = "Revenue is stored in the revenue column"
16sql_candidate = "SELECT region, name, revenue FROM customers ORDER BY revenue DESC LIMIT 5"
17execution_result = "5 rows returned"
18num_rows = 5
19
20# Format the prompt
21messages = [
22 {
23 "role": "system",
24 "content": "You are a judge that can check whether a given SQL correctly answers a given natural language user query. You'll be given Database Schema, Question, External Knowledge, SQL, logprob Score and its Execution Result.",
25 },
26 {
27 "role": "user",
28 "content": (
29 f"-- Database Schema: \n{db_schema}\n"
30 f"-- Question: {question}\n"
31 f"-- External Knowledge: {evidence}\n"
32 f"-- SQL: {sql_candidate}\n"
33 f"-- Execution Result #rows: {num_rows}\n"
34 f"-- Execution Result START\n{execution_result}\n"
35 f"-- END Execution Result\n"
36 f"-- Does SQL correctly answer Question?\n"
37 ),
38 }
39]
40
41# Generate reward score
42inputs = tokenizer.apply_chat_template(
43 messages,
44 tokenize=True,
45 add_generation_prompt=True,
46 return_tensors="pt",
47).to(model.device)
48
49outputs = model.generate(**inputs, max_new_tokens=10)
50score = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
51print(f"Reward score: {score}")1# Clone the repository
2git clone https://github.com/ContextualAI/bird-sql.git
3cd bird-sql
4
5# Install dependencies
6pip install -r requirements.txt
7
8# Download the reward model
9mkdir -p models/reward
10huggingface-cli download ContextualAI/ctx-bird-reward-250121 \
11 --local-dir models/reward1# Step 1: Generate SQL candidates
2python src/generate.py \
3 --input_file data/test_all.jsonl \
4 --output_dir output/generations/ \
5 --num_gpus 2
6
7# Step 2: Execute SQL candidates
8python src/process_sqls.py \
9 --input_file data/test_all.jsonl \
10 --generations_dir output/generations/ \
11 --output_dir output/with_results/ \
12 --compare_against_gt \
13 --sql_timeout 30.0
14
15# Step 3: Score with reward model
16VLLM_USE_V1=0 python src/reward.py \
17 --input_file output/with_results/data_with_results.jsonl \
18 --output_dir output/with_rewards \
19 --num_gpus 2
20
21# Step 4: Select best candidates
22python src/analysis.py \
23 --rewards_dir output/with_rewards \
24 --gt_sql_file data/test_gold_sqls.txt \
25 --output_dir output/analysis \
26 --num_cpus 100| Model | BIRD Dev Set | BIRD Test Set |
|---|---|---|
| Contextual-SQL | 73.50 | 75.63 |
1@misc{agrawal2025text2sql,
2 author = {Sheshansh Agrawal and Thien Nguyen},
3 title = {Open-Sourcing the Best Local Text-to-SQL System},
4 year = {2025},
5 url = {https://contextual.ai/blog/open-sourcing-the-best-local-text-to-sql-system/}
6}