Add src folder

src/ask.py

@@ -0,0 +1,93 @@
+"""
+TERMINAL CHAT WITH DATABASE
+Run:
+python src/ask.py chinook_1
+"""
+
+import sys
+from text2sql_engine import get_engine
+
+
+# -------------------------------
+# Pretty table printer
+# -------------------------------
+def print_table(cols, rows, limit=20):
+    if not rows or not cols:
+        print("No results\n")
+        return
+
+    cols = [str(c) for c in cols]
+
+    widths = [max(len(c), 12) for c in cols]
+
+    for r in rows[:limit]:
+        for i, val in enumerate(r):
+            widths[i] = max(widths[i], len(str(val)))
+
+    header = " | ".join(cols[i].ljust(widths[i]) for i in range(len(cols)))
+    print("\n" + header)
+    print("-" * len(header))
+
+    for r in rows[:limit]:
+        print(" | ".join(str(r[i]).ljust(widths[i]) for i in range(len(cols))))
+
+    if len(rows) > limit:
+        print(f"\n... showing first {limit} rows of {len(rows)}")
+
+    print()
+
+
+# -------------------------------
+# Main loop
+# -------------------------------
+def main():
+    if len(sys.argv) < 2:
+        print("Usage: python src/ask.py <db_id>")
+        return
+
+    db_id = sys.argv[1].strip()
+
+    print("Loading model... (first time takes 20-40s)")
+    engine = get_engine()
+
+    print(f"\nConnected to database: {db_id}")
+    print("Type 'exit' to quit\n")
+
+    while True:
+        try:
+            q = input("Ask> ").strip()
+
+            if not q:
+                continue
+
+            if q.lower() in ["exit", "quit"]:
+                break
+
+            result = engine.ask(q, db_id)
+
+            if result is None:
+                print("Model returned no output\n")
+                continue
+
+            print("\nGenerated SQL:")
+            print(result.get("sql", "<no sql>"))
+
+            if result.get("error"):
+                print("\nSQL Error:")
+                print(result["error"])
+            else:
+                print_table(
+                    result.get("columns", []),
+                    result.get("rows", []),
+                )
+
+        except KeyboardInterrupt:
+            break
+        except Exception as e:
+            print("\nRuntime error:", e, "\n")
+
+    print("\nBye!")
+
+
+if __name__ == "__main__":
+    main()

src/component_analysis.py

@@ -0,0 +1,229 @@
+import json
+import sqlite3
+import torch
+import re
+import matplotlib.pyplot as plt
+import numpy as np
+from pathlib import Path
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+DB_ROOT = PROJECT_ROOT / "data" / "database"
+
+# -------------------------------
+# Extract SQL components
+# -------------------------------
+def extract_components(sql):
+    sql = sql.lower()
+    return {
+        "select": "select" in sql,
+        "where": "where" in sql,
+        "group": "group by" in sql,
+        "order": "order by" in sql,
+        "and_or": (" and " in sql) or (" or " in sql),
+        "join": "join" in sql
+    }
+
+# -------------------------------
+# Fallback Difficulty Estimator
+# -------------------------------
+def estimate_difficulty(sql):
+    """Fallback if 'difficulty' is missing from the JSON."""
+    sql = sql.lower()
+    joins = sql.count("join")
+    conditions = sql.count("and") + sql.count("or")
+    
+    if "intersect" in sql or "except" in sql or "union" in sql or joins > 2:
+        return "extra"
+    elif joins == 2 or ("group by" in sql and conditions > 0):
+        return "hard"
+    elif joins == 1 or "group by" in sql or "order by" in sql:
+        return "medium"
+    else:
+        return "easy"
+
+# -------------------------------
+# Load schema
+# -------------------------------
+def load_schema(db_path):
+    conn = sqlite3.connect(db_path)
+    conn.text_factory = lambda b: b.decode(errors='ignore')
+    cursor = conn.cursor()
+
+    tables = cursor.execute(
+        "SELECT name FROM sqlite_master WHERE type='table';"
+    ).fetchall()
+
+    schema = ""
+    for (table,) in tables:
+        cols = cursor.execute(f"PRAGMA table_info({table});").fetchall()
+        col_names = [c[1] for c in cols]
+        schema += f"{table}({', '.join(col_names)})\n"
+
+    conn.close()
+    return schema
+
+# -------------------------------
+# Prompt
+# -------------------------------
+def build_prompt(question, schema):
+    return f"""Database Schema:
+{schema}
+
+Translate English to SQL:
+{question}
+SQL:
+"""
+
+# -------------------------------
+# Main
+# -------------------------------
+def main():
+    adapter = "checkpoints/rl_step_1800"
+    base_model = "Salesforce/codet5-base"
+
+    device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
+
+    print("Loading tokenizer and models...")
+    tokenizer = AutoTokenizer.from_pretrained(adapter)
+    base = AutoModelForSeq2SeqLM.from_pretrained(base_model).to(device)
+    model = PeftModel.from_pretrained(base, adapter).to(device)
+    model = model.merge_and_unload()
+    model.eval()
+
+    dev_json = PROJECT_ROOT / "data" / "dev.json"
+
+    with open(dev_json) as f:
+        dev = json.load(f)[:1000]  # Adjust number to test more/less
+
+    components_list = ["select", "where", "group", "order", "and_or", "join"]
+    difficulties_list = ["easy", "medium", "hard", "extra"]
+
+    # Nested dictionary for components
+    stats = {
+        comp: {diff: {"correct": 0, "total": 0} for diff in difficulties_list}
+        for comp in components_list
+    }
+
+    # 🚀 NEW: Trackers for OVERALL accuracy by difficulty
+    overall_correct = {diff: 0 for diff in difficulties_list}
+    overall_total = {diff: 0 for diff in difficulties_list}
+
+    print(f"\nRunning grouped evaluation on {len(dev)} examples...\n")
+
+    for i, ex in enumerate(dev, 1):
+        question = ex["question"]
+        gold_sql = ex["query"]
+        db_id = ex["db_id"]
+        
+        # Determine difficulty
+        difficulty = ex.get("difficulty", estimate_difficulty(gold_sql))
+        if difficulty not in difficulties_list:
+            difficulty = "medium"
+
+        db_path = DB_ROOT / db_id / f"{db_id}.sqlite"
+        schema = load_schema(db_path)
+        prompt = build_prompt(question, schema)
+
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=1000,
+                num_beams=4,
+                do_sample=False
+            )
+
+        pred_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        if "SQL:" in pred_sql:
+            pred_sql = pred_sql.split("SQL:")[-1]
+
+        # --- 1. Update Overall Accuracy Trackers ---
+        overall_total[difficulty] += 1
+        # Simple string match for quick overall accuracy
+        if pred_sql.strip().lower() == gold_sql.strip().lower():
+            overall_correct[difficulty] += 1
+
+        # --- 2. Update Component Stats ---
+        pred_comp = extract_components(pred_sql)
+        gold_comp = extract_components(gold_sql)
+
+        for comp in components_list:
+            if gold_comp[comp]:  # If the gold SQL required this component
+                stats[comp][difficulty]["total"] += 1
+                if pred_comp[comp]: # If the model successfully generated it
+                    stats[comp][difficulty]["correct"] += 1
+
+        if i % 20 == 0:
+            print(f"Processed {i}/{len(dev)}")
+
+    # -------------------------------
+    # Plotting (Grouped Bar Chart)
+    # -------------------------------
+    x = np.arange(len(components_list))
+    width = 0.2
+
+    def get_acc(diff):
+        return [
+            (stats[comp][diff]["correct"] / stats[comp][diff]["total"] * 100) if stats[comp][diff]["total"] > 0 else 0 
+            for comp in components_list
+        ]
+
+    acc_easy = get_acc("easy")
+    acc_medium = get_acc("medium")
+    acc_hard = get_acc("hard")
+    acc_extra = get_acc("extra")
+
+    fig, ax = plt.subplots(figsize=(14, 7))
+
+    bars1 = ax.bar(x - 1.5 * width, acc_easy, width, label='Easy', color='#2ecc71')
+    bars2 = ax.bar(x - 0.5 * width, acc_medium, width, label='Medium', color='#f1c40f')
+    bars3 = ax.bar(x + 0.5 * width, acc_hard, width, label='Hard', color='#e67e22')
+    bars4 = ax.bar(x + 1.5 * width, acc_extra, width, label='Extra', color='#e74c3c')
+
+    ax.set_ylabel('Accuracy (%)', fontsize=12)
+    ax.set_title('SQL Component Match Accuracy by Difficulty Level', fontsize=14, fontweight='bold')
+    ax.set_xticks(x)
+    ax.set_xticklabels([c.upper() for c in components_list], fontsize=11)
+    ax.legend(title="Query Difficulty")
+    ax.set_ylim(0, 115)
+
+    def autolabel(rects):
+        for rect in rects:
+            height = rect.get_height()
+            if height > 0:
+                ax.annotate(f'{int(height)}%',
+                            xy=(rect.get_x() + rect.get_width() / 2, height),
+                            xytext=(0, 3),
+                            textcoords="offset points",
+                            ha='center', va='bottom', fontsize=8, rotation=90)
+
+    autolabel(bars1)
+    autolabel(bars2)
+    autolabel(bars3)
+    autolabel(bars4)
+
+    ax.yaxis.grid(True, linestyle='--', alpha=0.7)
+    plt.tight_layout()
+    plt.savefig("component_by_difficulty_plot.png", dpi=300)
+
+    # -------------------------------
+    # 🚀 Terminal Printout
+    # -------------------------------
+    print("\n✅ Saved merged plot -> component_by_difficulty_plot.png")
+    
+    print("\n========================================")
+    print("🏆 OVERALL AVERAGE ACCURACY BY DIFFICULTY")
+    print("========================================")
+    for diff in difficulties_list:
+        if overall_total[diff] > 0:
+            avg = round((overall_correct[diff] / overall_total[diff]) * 100, 2)
+            print(f"{diff.capitalize():<8}: {avg:>5}%  ({overall_correct[diff]}/{overall_total[diff]} queries)")
+        else:
+            print(f"{diff.capitalize():<8}:   N/A  (0 queries)")
+    print("========================================\n")
+
+if __name__ == "__main__":
+    main()

src/convert_to_hf_dataset.py

@@ -0,0 +1,8 @@
+from datasets import Dataset
+import pandas as pd
+
+df = pd.read_csv("../data/processed/train.csv")
+ds = Dataset.from_pandas(df)
+ds.save_to_disk("../data/processed/train")
+print("DONE")
+

src/eval_baseline_codet5.py

@@ -0,0 +1,112 @@
+import json
+import sqlite3
+from pathlib import Path
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+# ---------------- PROMPT (same style as training) ----------------
+def build_prompt(question, schema):
+    return f"""translate English to SQL:
+
+Schema:
+{schema}
+
+Question:
+{question}
+
+SQL:"""
+
+# ---------------- LOAD SCHEMA ----------------
+def load_schema(db_path):
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    tables = cursor.execute(
+        "SELECT name FROM sqlite_master WHERE type='table';"
+    ).fetchall()
+
+    schema = ""
+    for (table,) in tables:
+        cols = cursor.execute(f"PRAGMA table_info({table});").fetchall()
+        col_names = [c[1] for c in cols]
+        schema += f"{table}({', '.join(col_names)})\n"
+
+    conn.close()
+    return schema
+
+# ---------------- EXECUTION MATCH ----------------
+def execution_match(pred_sql, gold_sql, db_path):
+    try:
+        conn = sqlite3.connect(db_path)
+        cur = conn.cursor()
+
+        cur.execute(pred_sql)
+        pred = cur.fetchall()
+
+        cur.execute(gold_sql)
+        gold = cur.fetchall()
+
+        conn.close()
+        return pred == gold
+
+    except Exception:
+        return False
+
+# ---------------- MAIN ----------------
+def main():
+    project_root = Path(__file__).resolve().parents[1]
+
+    dev_json = project_root / "data" / "dev.json"
+    db_root = project_root / "data" / "database"
+
+    device = "mps" if torch.backends.mps.is_available() else "cpu"
+
+    print("Loading BASE CodeT5...")
+    tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-base")
+    model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/codet5-base").to(device)
+    model.eval()
+
+    with open(dev_json) as f:
+        dev = json.load(f)[:100]
+
+    correct = 0
+
+    print(f"\nEvaluating {len(dev)} samples...\n")
+
+    for i, ex in enumerate(dev, 1):
+        question = ex["question"]
+        db_id = ex["db_id"]
+        gold_sql = ex["query"]
+
+        db_path = db_root / db_id / f"{db_id}.sqlite"
+        schema = load_schema(db_path)
+
+        prompt = build_prompt(question, schema)
+
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=80,
+                num_beams=4,
+                do_sample=False
+            )
+
+        pred_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        if "SQL:" in pred_sql:
+            pred_sql = pred_sql.split("SQL:")[-1].strip()
+
+        if execution_match(pred_sql, gold_sql, db_path):
+            correct += 1
+
+        if i % 10 == 0:
+            print(f"{i}/100 | Accuracy: {correct/i:.3f}")
+
+    print("\n=============================")
+    print(f"BASE MODEL ACCURACY: {correct}% / 100 = {correct}%")
+    print("=============================")
+
+if __name__ == "__main__":
+    main()

src/eval_both_metrics.py

@@ -0,0 +1,144 @@
+import json
+import sqlite3
+import torch
+import re
+import time
+import argparse
+from pathlib import Path
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+DB_ROOT = PROJECT_ROOT / "data" / "database"
+
+# -------------------------------
+# 1. NORMALIZATION FOR EXACT MATCH
+# -------------------------------
+def normalize_sql(sql):
+    """Cleans SQL to make Exact Match grading fair (ignores spacing/cases)."""
+    sql = sql.replace('"', "'")        # Standardize quotes
+    sql = re.sub(r"\s+", " ", sql)     # Remove extra spaces/newlines
+    sql = sql.strip().lower()          # Lowercase everything
+    sql = sql.rstrip(";")              # Remove trailing semicolons
+    return sql
+
+# -------------------------------
+# 2. EXECUTION ACCURACY CHECK
+# -------------------------------
+def check_execution(pred_sql, gold_sql, db_path):
+    """Runs both queries and checks if the output rows/columns match."""
+    try:
+        conn = sqlite3.connect(db_path)
+        # Handle bad characters in Spider DBs
+        conn.text_factory = lambda b: b.decode(errors='ignore')
+        
+        # 5-second timeout
+        start_time = time.monotonic()
+        def timeout_handler():
+            return 1 if (time.monotonic() - start_time) > 5.0 else 0
+        conn.set_progress_handler(timeout_handler, 10000)
+
+        cursor = conn.cursor()
+
+        # Get Predicted Result
+        cursor.execute(pred_sql)
+        pred_res = cursor.fetchall()
+
+        # Get Gold Result
+        cursor.execute(gold_sql)
+        gold_res = cursor.fetchall()
+
+        conn.close()
+        return pred_res == gold_res
+    except Exception:
+        return False
+
+# -------------------------------
+# 3. LOAD SCHEMA
+# -------------------------------
+def load_schema(db_path):
+    conn = sqlite3.connect(db_path)
+    conn.text_factory = lambda b: b.decode(errors='ignore')
+    cursor = conn.cursor()
+    tables = cursor.execute("SELECT name FROM sqlite_master WHERE type='table';").fetchall()
+    schema = ""
+    for (table,) in tables:
+        cols = cursor.execute(f"PRAGMA table_info({table});").fetchall()
+        col_names = [c[1] for c in cols]
+        schema += f"{table}({', '.join(col_names)})\n"
+    conn.close()
+    return schema
+
+# -------------------------------
+# 4. MAIN PIPELINE
+# -------------------------------
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--adapter", type=str, required=True, help="Path to your SFT or RLHF checkpoint")
+    parser.add_argument("--num_samples", type=int, default=1034, help="How many samples to evaluate")
+    args = parser.parse_args()
+
+    device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
+    base_model = "Salesforce/codet5-base"
+
+    print(f"\n🚀 Loading Model from: {args.adapter}")
+    tokenizer = AutoTokenizer.from_pretrained(args.adapter)
+    base = AutoModelForSeq2SeqLM.from_pretrained(base_model).to(device)
+    model = PeftModel.from_pretrained(base, args.adapter).to(device)
+    model = model.merge_and_unload()
+    model.eval()
+
+    dev_json = PROJECT_ROOT / "data" / "dev.json"
+    with open(dev_json) as f:
+        dev = json.load(f)[:args.num_samples]
+
+    em_correct = 0
+    ex_correct = 0
+    total = len(dev)
+
+    print(f"\n📊 Evaluating {total} queries for BOTH Exact Match and Execution Accuracy...\n")
+
+    for i, ex in enumerate(dev, 1):
+        question = ex["question"]
+        gold_sql = ex["query"]
+        db_id = ex["db_id"]
+        db_path = DB_ROOT / db_id / f"{db_id}.sqlite"
+
+        # Generate SQL
+        schema = load_schema(db_path)
+        prompt = f"Database Schema:\n{schema}\nTranslate English to SQL:\n{question}\nSQL:\n"
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+        with torch.no_grad():
+            outputs = model.generate(**inputs, max_new_tokens=100, num_beams=4, do_sample=False)
+        
+        pred_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        if "SQL:" in pred_sql:
+            pred_sql = pred_sql.split("SQL:")[-1].strip()
+
+        # --- METRIC 1: EXACT MATCH ---
+        is_em = (normalize_sql(pred_sql) == normalize_sql(gold_sql))
+        if is_em:
+            em_correct += 1
+
+        # --- METRIC 2: EXECUTION ACCURACY ---
+        is_ex = check_execution(pred_sql, gold_sql, db_path)
+        if is_ex:
+            ex_correct += 1
+
+        if i % 50 == 0 or i == total:
+            print(f"Progress: {i}/{total} | Current EM: {(em_correct/i)*100:.2f}% | Current EX: {(ex_correct/i)*100:.2f}%")
+
+    # Final Results
+    final_em = (em_correct / total) * 100
+    final_ex = (ex_correct / total) * 100
+
+    print("\n==========================================")
+    print(f"🎯 FINAL RESULTS FOR: {args.adapter}")
+    print("==========================================")
+    print(f"Exact Match (EM) Accuracy      : {final_em:.2f}%")
+    print(f"Execution (EX) Accuracy        : {final_ex:.2f}%")
+    print("==========================================\n")
+
+if __name__ == "__main__":
+    main()

src/eval_rl_fixed.py

@@ -0,0 +1,466 @@
+# import json
+# import sqlite3
+# import argparse
+# from pathlib import Path
+# import torch
+# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+# from peft import PeftModel
+
+# # ---------------- PROMPT (IDENTICAL TO TRAINING) ----------------
+# def build_prompt(question, schema):
+#     return f"""
+# Database Schema:
+# {schema}
+
+# Translate English to SQL:
+# {question}
+# SQL:
+# """
+
+# # ---------------- LOAD SCHEMA ----------------
+# def load_schema(db_path):
+#     conn = sqlite3.connect(db_path)
+#     cursor = conn.cursor()
+
+#     tables = cursor.execute(
+#         "SELECT name FROM sqlite_master WHERE type='table';"
+#     ).fetchall()
+
+#     schema = ""
+#     for (table,) in tables:
+#         cols = cursor.execute(f"PRAGMA table_info({table});").fetchall()
+#         col_names = [c[1] for c in cols]
+#         schema += f"{table}({', '.join(col_names)})\n"
+
+#     conn.close()
+#     return schema
+
+
+# # ---------------- EXECUTION CHECK ----------------
+# def execution_match(pred_sql, gold_sql, db_path):
+#     try:
+#         conn = sqlite3.connect(db_path)
+#         cur = conn.cursor()
+
+#         cur.execute(pred_sql)
+#         pred = cur.fetchall()
+
+#         cur.execute(gold_sql)
+#         gold = cur.fetchall()
+
+#         conn.close()
+#         return pred == gold
+
+#     except Exception:
+#         return False
+
+
+# # ---------------- MAIN ----------------
+# def main():
+#     parser = argparse.ArgumentParser()
+#     parser.add_argument("--adapter", type=str, required=True)
+#     parser.add_argument("--num_samples", type=int, default=1034)
+#     args = parser.parse_args()
+
+#     project_root = Path(__file__).resolve().parents[1]
+
+#     dev_json = project_root / "data" / "dev.json"
+#     db_root = project_root / "data" / "database"
+
+#     device = "mps" if torch.backends.mps.is_available() else "cpu"
+
+#     # load model
+#     base_model = "Salesforce/codet5-base"
+#     tokenizer = AutoTokenizer.from_pretrained(args.adapter)
+#     base = AutoModelForSeq2SeqLM.from_pretrained(base_model).to(device)
+#     model = PeftModel.from_pretrained(base, args.adapter).to(device)
+#     model = model.merge_and_unload()
+
+#     with open(dev_json) as f:
+#         dev = json.load(f)[: args.num_samples]
+
+#     correct = 0
+
+#     print(f"Evaluating {len(dev)} examples...\n")
+
+#     for i, ex in enumerate(dev, 1):
+#         question = ex["question"]
+#         db_id = ex["db_id"]
+#         gold_sql = ex["query"]
+
+#         db_path = db_root / db_id / f"{db_id}.sqlite"
+#         schema = load_schema(db_path)
+
+#         prompt = build_prompt(question, schema)
+
+#         inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+#         with torch.no_grad():
+#             outputs = model.generate(
+#                 **inputs,
+#                 max_new_tokens=80,
+#                 do_sample=False,
+#                 num_beams=4,
+#             )
+
+#         pred_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+#         if "SQL:" in pred_sql:
+#             pred_sql = pred_sql.split("SQL:")[-1].strip()
+
+#         match = execution_match(pred_sql, gold_sql, db_path)
+
+#         if match:
+#             correct += 1
+
+#         if i % 10 == 0:
+#             print(f"{i}/{len(dev)} | Acc: {correct/i:.3f}")
+
+#     print("\n=============================")
+#     print(f"FINAL EXECUTION ACCURACY: {correct/len(dev)*100:.2f}%")
+#     print("=============================")
+
+
+# if __name__ == "__main__":
+#     main()
+
+
+# import json
+# import sqlite3
+# import argparse
+# import time
+# from pathlib import Path
+# import torch
+# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+# from peft import PeftModel
+
+# # ---------------- PROMPT (IDENTICAL TO TRAINING) ----------------
+# def build_prompt(question, schema):
+#     return f"""
+# Database Schema:
+# {schema}
+
+# Translate English to SQL:
+# {question}
+# SQL:
+# """
+
+# # ---------------- LOAD SCHEMA ----------------
+# def load_schema(db_path):
+#     conn = sqlite3.connect(db_path)
+#     cursor = conn.cursor()
+
+#     tables = cursor.execute(
+#         "SELECT name FROM sqlite_master WHERE type='table';"
+#     ).fetchall()
+
+#     schema = ""
+#     for (table,) in tables:
+#         cols = cursor.execute(f"PRAGMA table_info({table});").fetchall()
+#         col_names = [c[1] for c in cols]
+#         schema += f"{table}({', '.join(col_names)})\n"
+
+#     conn.close()
+#     return schema
+
+
+# # ---------------- EXECUTION CHECK WITH TIMEOUT ----------------
+# def execution_match(pred_sql, gold_sql, db_path):
+#     try:
+#         conn = sqlite3.connect(db_path)
+        
+#         # --- 5-SECOND TIMEOUT SO EVALUATION DOESN'T FREEZE ---
+#         start_time = time.monotonic()
+#         def timeout_handler():
+#             return 1 if (time.monotonic() - start_time) > 5.0 else 0
+#         conn.set_progress_handler(timeout_handler, 10000)
+
+#         cur = conn.cursor()
+
+#         cur.execute(pred_sql)
+#         pred = cur.fetchall()
+
+#         cur.execute(gold_sql)
+#         gold = cur.fetchall()
+
+#         conn.close()
+#         return pred == gold
+
+#     except Exception:
+#         return False
+
+
+# # ---------------- MAIN ----------------
+# def main():
+#     parser = argparse.ArgumentParser()
+#     parser.add_argument("--adapter", type=str, required=True)
+#     parser.add_argument("--num_samples", type=int, default=1034)
+#     args = parser.parse_args()
+
+#     project_root = Path(__file__).resolve().parents[1]
+
+#     dev_json = project_root / "data" / "dev.json"
+#     db_root = project_root / "data" / "database"
+
+#     # 🎯 Added CUDA support for Nvidia GPUs
+#     device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
+
+#     # load model
+#     base_model = "Salesforce/codet5-base"
+#     print(f"Loading Base: {base_model}")
+#     print(f"Loading Adapter: {args.adapter}")
+    
+#     tokenizer = AutoTokenizer.from_pretrained(args.adapter)
+#     base = AutoModelForSeq2SeqLM.from_pretrained(base_model).to(device)
+#     model = PeftModel.from_pretrained(base, args.adapter).to(device)
+#     model = model.merge_and_unload()
+
+#     with open(dev_json) as f:
+#         dev = json.load(f)[: args.num_samples]
+
+#     correct = 0
+
+#     print(f"Evaluating {len(dev)} examples...\n")
+
+#     for i, ex in enumerate(dev, 1):
+#         question = ex["question"]
+#         db_id = ex["db_id"]
+#         gold_sql = ex["query"]
+
+#         db_path = db_root / db_id / f"{db_id}.sqlite"
+#         schema = load_schema(db_path)
+
+#         prompt = build_prompt(question, schema)
+
+#         inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+#         with torch.no_grad():
+#             outputs = model.generate(
+#                 **inputs,
+#                 max_new_tokens=80,
+#                 do_sample=False,
+#                 num_beams=4,
+#             )
+
+#         pred_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+#         if "SQL:" in pred_sql:
+#             pred_sql = pred_sql.split("SQL:")[-1].strip()
+
+#         match = execution_match(pred_sql, gold_sql, db_path)
+
+#         if match:
+#             correct += 1
+
+#         if i % 10 == 0:
+#             print(f"{i}/{len(dev)} | Acc: {correct/i:.3f}")
+
+#     print("\n=============================")
+#     print(f"FINAL EXECUTION ACCURACY: {correct/len(dev)*100:.2f}%")
+#     print("=============================")
+
+
+# if __name__ == "__main__":
+#     main()
+
+
+import json
+import subprocess
+import sys
+import argparse
+import random
+import sqlite3
+import time
+import re
+from pathlib import Path
+
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+
+# Assuming you have a prompting.py that has encode_prompt
+from prompting import encode_prompt
+
+# -------------------------------
+# LIVE CHECK HELPERS
+# -------------------------------
+def normalize_sql(sql):
+    """Basic normalization for the live progress bar."""
+    sql = sql.replace('"', "'")
+    sql = re.sub(r"\s+", " ", sql)
+    return sql.strip().lower().rstrip(";")
+
+def check_execution(pred_sql, gold_sql, db_path):
+    """Basic execution check for the live progress bar."""
+    try:
+        conn = sqlite3.connect(db_path)
+        conn.text_factory = lambda b: b.decode(errors='ignore')
+        
+        # 2-second timeout so the live tracker doesn't freeze forever
+        start_time = time.monotonic()
+        def timeout_handler():
+            return 1 if (time.monotonic() - start_time) > 2.0 else 0
+        conn.set_progress_handler(timeout_handler, 10000)
+
+        cursor = conn.cursor()
+        cursor.execute(pred_sql)
+        pred_res = cursor.fetchall()
+        
+        cursor.execute(gold_sql)
+        gold_res = cursor.fetchall()
+        conn.close()
+        
+        # Simple sorted check for the live tracker
+        return sorted(pred_res) == sorted(gold_res)
+    except Exception:
+        return False
+
+# -------------------------------
+# SPIDER PARSER
+# -------------------------------
+def _parse_spider_accuracy(stdout: str, metric_type: str) -> float | None:
+    for line in stdout.splitlines():
+        if metric_type == "exec" and line.strip().startswith("execution"):
+            try: return float(line.split()[-1])
+            except: pass
+        elif metric_type == "match" and line.strip().startswith("exact"):
+            try: return float(line.split()[-1])
+            except: pass
+    return None
+
+# -------------------------------
+# MAIN
+# -------------------------------
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--adapter", type=str, required=True, help="Path to your SFT or RLHF checkpoint")
+    parser.add_argument("--num_samples", type=int, default=700, help="Number of samples to evaluate")
+    parser.add_argument("--shuffle_dev", action="store_true")
+    parser.add_argument("--shuffle_seed", type=int, default=42)
+    args = parser.parse_args()
+
+    project_root = Path(__file__).resolve().parents[1]
+    adapter_dir = project_root / args.adapter
+
+    db_root = project_root / "data" / "database"
+    table_json = project_root / "data" / "tables.json"
+    dev_json = project_root / "data" / "dev.json"
+    
+    pred_path = project_root / "temp_predictions.txt"
+    temp_gold_path = project_root / "temp_gold.sql"
+
+    if not adapter_dir.exists():
+        raise FileNotFoundError(f"Missing adapter dir: {adapter_dir}")
+
+    device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    BASE_MODEL = "Salesforce/codet5-base"
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    print(f"Loading Model: {args.adapter}...")
+    base = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL).to(device)
+    model = PeftModel.from_pretrained(base, str(adapter_dir)).to(device)
+    model = model.merge_and_unload()
+    model.eval()
+
+    with dev_json.open() as f:
+        dev = json.load(f)
+
+    if args.shuffle_dev:
+        rng = random.Random(args.shuffle_seed)
+        rng.shuffle(dev)
+
+    dev = dev[: args.num_samples]
+    total = len(dev)
+
+    gen_kwargs = dict(
+        max_new_tokens=160,
+        num_beams=4,
+        do_sample=False,
+        early_stopping=True,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+
+    print(f"\n🚀 Generating and live-tracking {total} samples...\n")
+
+    em_correct = 0
+    ex_correct = 0
+
+    with pred_path.open("w") as out_pred, temp_gold_path.open("w") as out_gold, torch.no_grad():
+        for i, ex in enumerate(dev, start=1):
+            db_id = ex["db_id"]
+            question = ex["question"]
+            gold_query = ex["query"]
+            db_path = db_root / db_id / f"{db_id}.sqlite"
+
+            # Generate
+            input_ids = encode_prompt(tokenizer, question, db_id, device=device, max_input_tokens=512)
+            input_ids = input_ids.unsqueeze(0).to(device)
+            attention_mask = (input_ids != tokenizer.pad_token_id).long().to(device)
+
+            outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask, **gen_kwargs)
+            pred_sql = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+            
+            # Write to files for official spider eval later
+            out_pred.write(f"{pred_sql}\n")
+            out_gold.write(f"{gold_query}\t{db_id}\n")
+
+            # --- LIVE TRACKING CHECKS ---
+            if normalize_sql(pred_sql) == normalize_sql(gold_query):
+                em_correct += 1
+            if check_execution(pred_sql, gold_query, db_path):
+                ex_correct += 1
+
+            # Print progress every 50 loops
+            if i % 10 == 0 or i == total:
+                print(f"Progress: {i}/{total} | Current EM: {(em_correct/i)*100:.2f}% | Current EX: {(ex_correct/i)*100:.2f}%")
+
+    print("\nGeneration finished. Running Official Spider Evaluations for final numbers...\n")
+
+    eval_script = project_root / "spider_eval" / "evaluation.py"
+
+    # 1. RUN EXACT MATCH EVAL
+    cmd_match = [
+        sys.executable, str(eval_script),
+        "--gold", str(temp_gold_path),
+        "--pred", str(pred_path),
+        "--etype", "match",
+        "--db", str(db_root),
+        "--table", str(table_json),
+    ]
+    proc_match = subprocess.run(cmd_match, capture_output=True, text=True)
+    exact_acc = _parse_spider_accuracy(proc_match.stdout, "match")
+
+    # 2. RUN EXECUTION EVAL
+    cmd_exec = [
+        sys.executable, str(eval_script),
+        "--gold", str(temp_gold_path),
+        "--pred", str(pred_path),
+        "--etype", "exec",
+        "--db", str(db_root),
+        "--table", str(table_json),
+    ]
+    proc_exec = subprocess.run(cmd_exec, capture_output=True, text=True)
+    exec_acc = _parse_spider_accuracy(proc_exec.stdout, "exec")
+
+    print("==========================================")
+    print(f"🎯 OFFICIAL SPIDER RESULTS FOR: {args.adapter}")
+    print("==========================================")
+    
+    if exact_acc is not None:
+        print(f"Exact Set Match Accuracy  : {exact_acc*100:.2f}%")
+    else:
+        print("Exact Set Match Accuracy  : Could not parse output")
+        
+    if exec_acc is not None:
+        print(f"Execution Accuracy        : {exec_acc*100:.2f}%")
+    else:
+        print("Execution Accuracy        : Could not parse output")
+    print("==========================================\n")
+
+if __name__ == "__main__":
+    main()

src/eval_rl_t5.py

@@ -0,0 +1,279 @@
+# import sys
+# import os
+# sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+# import json
+
+# import subprocess
+
+# import argparse
+# from pathlib import Path
+
+# import torch
+# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+# from peft import PeftModel
+
+# # IMPORTANT: must match training prompt format
+# from prompting import build_prompt
+# from schema_utils import get_schema as get_db_schema
+
+
+# def _parse_exec_accuracy(stdout: str):
+#     for line in stdout.splitlines():
+#         if line.strip().startswith("execution"):
+#             parts = line.split()
+#             try:
+#                 return float(parts[-1])
+#             except Exception:
+#                 return None
+#     return None
+
+
+# def main():
+#     parser = argparse.ArgumentParser()
+#     parser.add_argument("--adapter", type=str, default="checkpoints/best_rlhf_model")
+#     parser.add_argument("--num_samples", type=int, default=200)
+#     args = parser.parse_args()
+
+#     project_root = Path(__file__).resolve().parents[1]
+#     adapter_dir = project_root / args.adapter
+
+#     if not adapter_dir.exists():
+#         raise FileNotFoundError(f"Adapter not found: {adapter_dir}")
+
+#     db_root = project_root / "data" / "database"
+#     table_json = project_root / "data" / "tables.json"
+#     dev_json = project_root / "data" / "dev.json"
+#     gold_sql = project_root / "data" / "dev_gold.sql"
+#     pred_path = project_root / "predictions_rl.txt"
+
+#     device = "mps" if torch.backends.mps.is_available() else "cpu"
+
+#     # ---- LOAD MODEL (CodeT5 + LoRA) ----
+#     base_model = "Salesforce/codet5-base"
+
+#     tokenizer = AutoTokenizer.from_pretrained(str(adapter_dir))
+#     base = AutoModelForSeq2SeqLM.from_pretrained(base_model).to(device)
+#     model = PeftModel.from_pretrained(base, str(adapter_dir)).to(device)
+
+#     # merge LoRA for faster inference
+#     model = model.merge_and_unload()
+#     model.eval()
+#     model.config.use_cache = True
+
+#     if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
+#         tokenizer.pad_token = tokenizer.eos_token
+
+#     # ---- LOAD DATA ----
+#     with dev_json.open() as f:
+#         dev = json.load(f)
+
+#     dev = dev[: args.num_samples]
+
+#     gen_kwargs = dict(
+#         max_new_tokens=120,
+#         do_sample=False,
+#         num_beams=1,
+#         pad_token_id=tokenizer.pad_token_id,
+#         eos_token_id=tokenizer.eos_token_id,
+#     )
+
+#     print(f"Generating {len(dev)} predictions...")
+
+#     with pred_path.open("w") as out_f, torch.no_grad():
+#         for i, ex in enumerate(dev, start=1):
+#             db_id = ex["db_id"]
+#             question = ex["question"]
+
+#             db_path = db_root / db_id / f"{db_id}.sqlite"
+#             schema = get_db_schema(str(db_path))
+#             prompt = build_prompt(question, schema, use_schema=True)
+
+#             inputs = tokenizer(
+#                 prompt,
+#                 return_tensors="pt",
+#                 truncation=True,
+#                 max_length=512
+#             ).to(device)
+
+#             out = model.generate(**inputs, **gen_kwargs)
+#             pred_sql = tokenizer.decode(out[0], skip_special_tokens=True).strip()
+
+#             out_f.write(f"{pred_sql}\t{db_id}\n")
+
+#             if i % 20 == 0 or i == len(dev):
+#                 print(f"{i}/{len(dev)} done")
+
+#     # ---- SPIDER OFFICIAL EVAL ----
+#     eval_script = project_root / "spider_eval" / "evaluation.py"
+
+#     cmd = [
+#         sys.executable,
+#         str(eval_script),
+#         "--gold",
+#         str(gold_sql),
+#         "--pred",
+#         str(pred_path),
+#         "--etype",
+#         "exec",
+#         "--db",
+#         str(db_root),
+#         "--table",
+#         str(table_json),
+#     ]
+
+#     print("\nRunning Spider execution evaluation...\n")
+#     proc = subprocess.run(cmd, capture_output=True, text=True)
+
+#     if proc.returncode != 0:
+#         print(proc.stdout)
+#         print(proc.stderr)
+#         sys.exit(proc.returncode)
+
+#     print(proc.stdout)
+
+#     acc = _parse_exec_accuracy(proc.stdout)
+#     if acc is not None:
+#         print(f"\nFINAL EXECUTION ACCURACY: {acc*100:.2f}%")
+#     else:
+#         print("Could not parse execution accuracy")
+
+
+# if __name__ == "__main__":
+#     main()
+
+
+import json
+import sqlite3
+import argparse
+import time
+from pathlib import Path
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+
+# ---------------- PROMPT (FIXED TO PERFECTLY MATCH RLHF TRAINING) ----------------
+def build_prompt(question, schema):
+    return f"translate English to SQL:\n\nSchema:\n{schema}\n\nQuestion:\n{question}\n\nSQL:"
+
+# ---------------- LOAD SCHEMA (FIXED TO MATCH TRAINING FORMAT) ----------------
+def load_schema(db_path):
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    tables = cursor.execute(
+        "SELECT name FROM sqlite_master WHERE type='table';"
+    ).fetchall()
+
+    schema = ""
+    for (table,) in tables:
+        cols = cursor.execute(f"PRAGMA table_info({table});").fetchall()
+        col_names = [c[1] for c in cols]
+        # Space-separated, not newline-separated, just like the RLHF script
+        schema += f"{table}({', '.join(col_names)}) " 
+
+    conn.close()
+    return schema.strip()
+
+
+# ---------------- EXECUTION CHECK WITH TIMEOUT ----------------
+def execution_match(pred_sql, gold_sql, db_path):
+    try:
+        conn = sqlite3.connect(db_path)
+        
+        # --- 5-SECOND TIMEOUT SO THE SCRIPT DOESN'T HANG ---
+        start_time = time.monotonic()
+        def timeout_handler():
+            return 1 if (time.monotonic() - start_time) > 5.0 else 0
+        conn.set_progress_handler(timeout_handler, 10000)
+
+        cur = conn.cursor()
+
+        cur.execute(pred_sql)
+        pred = cur.fetchall()
+
+        cur.execute(gold_sql)
+        gold = cur.fetchall()
+
+        conn.close()
+        return pred == gold
+
+    except Exception:
+        return False
+
+
+# ---------------- MAIN ----------------
+def main():
+    parser = argparse.ArgumentParser()
+    # 🎯 Set the default directly to your best RLHF model!
+    parser.add_argument("--adapter", type=str, default="checkpoints/rlhf_t5_best")
+    parser.add_argument("--num_samples", type=int, default=1000)
+    args = parser.parse_args()
+
+    project_root = Path(__file__).resolve().parents[1]
+    
+    # Resolve adapter path safely
+    adapter_path = project_root / args.adapter
+
+    dev_json = project_root / "data" / "dev.json"
+    db_root = project_root / "data" / "database"
+
+    # 🎯 Added CUDA support
+    device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
+
+    # load model
+    base_model = "t5-small"
+    print(f"Loading Base: {base_model}")
+    print(f"Loading Adapter: {adapter_path}")
+    
+    tokenizer = AutoTokenizer.from_pretrained(str(adapter_path))
+    base = AutoModelForSeq2SeqLM.from_pretrained(base_model).to(device)
+    model = PeftModel.from_pretrained(base, str(adapter_path)).to(device)
+    model = model.merge_and_unload()
+
+    with open(dev_json) as f:
+        dev = json.load(f)[: args.num_samples]
+
+    correct = 0
+
+    print(f"Evaluating {len(dev)} examples...\n")
+
+    for i, ex in enumerate(dev, 1):
+        question = ex["question"]
+        db_id = ex["db_id"]
+        gold_sql = ex["query"]
+
+        db_path = db_root / db_id / f"{db_id}.sqlite"
+        schema = load_schema(db_path)
+
+        prompt = build_prompt(question, schema)
+
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=80,
+                do_sample=False,
+                num_beams=4,
+            )
+
+        pred_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        if "SQL:" in pred_sql:
+            pred_sql = pred_sql.split("SQL:")[-1].strip()
+
+        match = execution_match(pred_sql, gold_sql, db_path)
+
+        if match:
+            correct += 1
+
+        if i % 10 == 0:
+            print(f"{i}/{len(dev)} | Acc: {correct/i:.3f}")
+
+    print("\n=============================")
+    print(f"FINAL EXECUTION ACCURACY: {correct/len(dev)*100:.2f}%")
+    print("=============================")
+
+
+if __name__ == "__main__":
+    main()

src/eval_single_model.py

@@ -0,0 +1,218 @@
+import json
+import subprocess
+import sys
+import argparse
+import random
+import sqlite3
+import time
+import re
+import matplotlib.pyplot as plt
+import numpy as np
+from pathlib import Path
+
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+
+# Assuming you have a prompting.py that has encode_prompt
+from prompting import encode_prompt
+
+# -------------------------------
+# LIVE CHECK HELPERS
+# -------------------------------
+def normalize_sql(sql):
+    sql = sql.replace('"', "'")
+    sql = re.sub(r"\s+", " ", sql)
+    return sql.strip().lower().rstrip(";")
+
+def check_execution(pred_sql, gold_sql, db_path):
+    try:
+        conn = sqlite3.connect(db_path)
+        conn.text_factory = lambda b: b.decode(errors='ignore')
+        
+        start_time = time.monotonic()
+        def timeout_handler():
+            return 1 if (time.monotonic() - start_time) > 2.0 else 0
+        conn.set_progress_handler(timeout_handler, 10000)
+
+        cursor = conn.cursor()
+        cursor.execute(pred_sql)
+        pred_res = cursor.fetchall()
+        
+        cursor.execute(gold_sql)
+        gold_res = cursor.fetchall()
+        conn.close()
+        
+        return sorted(pred_res) == sorted(gold_res)
+    except Exception:
+        return False
+
+# -------------------------------
+# SPIDER PARSER
+# -------------------------------
+def _parse_spider_accuracy(stdout: str, metric_type: str) -> float | None:
+    for line in stdout.splitlines():
+        if metric_type == "exec" and line.strip().startswith("execution"):
+            try: return float(line.split()[-1])
+            except: pass
+        elif metric_type == "match" and line.strip().startswith("exact"):
+            try: return float(line.split()[-1])
+            except: pass
+    return None
+
+# -------------------------------
+# MAIN
+# -------------------------------
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--adapter", type=str, required=True, help="Path to your checkpoint")
+    parser.add_argument("--base_model", type=str, required=True, help="E.g., facebook/bart-base, t5-small")
+    parser.add_argument("--model_name", type=str, required=True, help="Name for the plot label (e.g., 'BART RLHF')")
+    parser.add_argument("--num_samples", type=int, default=700)
+    args = parser.parse_args()
+
+    project_root = Path(__file__).resolve().parents[1]
+    adapter_dir = project_root / args.adapter
+
+    db_root = project_root / "data" / "database"
+    table_json = project_root / "data" / "tables.json"
+    dev_json = project_root / "data" / "dev.json"
+    
+    pred_path = project_root / "temp_predictions.txt"
+    temp_gold_path = project_root / "temp_gold.sql"
+    
+    # NEW: Plot directory setup
+    plot_dir = project_root / "comparison_plots"
+    plot_dir.mkdir(parents=True, exist_ok=True)
+    results_json_path = plot_dir / "all_metrics.json"
+
+    if not adapter_dir.exists():
+        raise FileNotFoundError(f"Missing adapter dir: {adapter_dir}")
+
+    device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Loading Base Model: {args.base_model} on {device}...")
+    
+    tokenizer = AutoTokenizer.from_pretrained(args.base_model)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    base = AutoModelForSeq2SeqLM.from_pretrained(args.base_model).to(device)
+    model = PeftModel.from_pretrained(base, str(adapter_dir)).to(device)
+    model = model.merge_and_unload()
+    model.eval()
+
+    with dev_json.open() as f:
+        dev = json.load(f)[: args.num_samples]
+    total = len(dev)
+
+    gen_kwargs = dict(
+        max_new_tokens=160,
+        num_beams=4,
+        do_sample=False,
+        early_stopping=True,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+
+    print(f"\n🚀 Generating and live-tracking {total} samples...\n")
+
+    em_correct = 0
+    ex_correct = 0
+
+    with pred_path.open("w") as out_pred, temp_gold_path.open("w") as out_gold, torch.no_grad():
+        for i, ex in enumerate(dev, start=1):
+            db_id = ex["db_id"]
+            question = ex["question"]
+            gold_query = ex["query"]
+            db_path = db_root / db_id / f"{db_id}.sqlite"
+
+            # Generate
+            input_ids = encode_prompt(tokenizer, question, db_id, device=device, max_input_tokens=512)
+            input_ids = input_ids.unsqueeze(0).to(device)
+            attention_mask = (input_ids != tokenizer.pad_token_id).long().to(device)
+
+            outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask, **gen_kwargs)
+            pred_sql = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+            
+            out_pred.write(f"{pred_sql}\n")
+            out_gold.write(f"{gold_query}\t{db_id}\n")
+
+            # --- PRINT FIRST 3 EXAMPLES ---
+            if i <= 3:
+                print(f"--- 🔍 Example {i} ---")
+                print(f"Q   : {question}")
+                print(f"Gold: {gold_query}")
+                print(f"Pred: {pred_sql}")
+                print("-" * 25)
+
+            # --- LIVE TRACKING CHECKS ---
+            if normalize_sql(pred_sql) == normalize_sql(gold_query):
+                em_correct += 1
+            if check_execution(pred_sql, gold_query, db_path):
+                ex_correct += 1
+
+            if i % 50 == 0 or i == total:
+                print(f"Progress: {i}/{total} | Current EM: {(em_correct/i)*100:.2f}% | Current EX: {(ex_correct/i)*100:.2f}%")
+
+    print("\nRunning Official Spider Evaluations...")
+    eval_script = project_root / "spider_eval" / "evaluation.py"
+
+    proc_match = subprocess.run([sys.executable, str(eval_script), "--gold", str(temp_gold_path), "--pred", str(pred_path), "--etype", "match", "--db", str(db_root), "--table", str(table_json)], capture_output=True, text=True)
+    exact_acc = _parse_spider_accuracy(proc_match.stdout, "match")
+
+    proc_exec = subprocess.run([sys.executable, str(eval_script), "--gold", str(temp_gold_path), "--pred", str(pred_path), "--etype", "exec", "--db", str(db_root), "--table", str(table_json)], capture_output=True, text=True)
+    exec_acc = _parse_spider_accuracy(proc_exec.stdout, "exec")
+
+    print("\n==========================================")
+    print(f"🎯 RESULTS FOR: {args.model_name}")
+    print("==========================================")
+    exact_val = exact_acc * 100 if exact_acc else 0
+    exec_val = exec_acc * 100 if exec_acc else 0
+    print(f"Exact Match : {exact_val:.2f}%")
+    print(f"Execution   : {exec_val:.2f}%")
+    print("==========================================\n")
+
+    # -------------------------------
+    # SAVE JSON & GENERATE PLOT
+    # -------------------------------
+    if results_json_path.exists():
+        with open(results_json_path, 'r') as f:
+            all_results = json.load(f)
+    else:
+        all_results = {}
+
+    all_results[args.model_name] = {"EM": exact_val, "EX": exec_val}
+
+    with open(results_json_path, 'w') as f:
+        json.dump(all_results, f, indent=4)
+
+    labels = list(all_results.keys())
+    em_vals = [all_results[k]["EM"] for k in labels]
+    ex_vals = [all_results[k]["EX"] for k in labels]
+
+    x = np.arange(len(labels))
+    width = 0.35
+
+    plt.figure(figsize=(max(8, len(labels) * 1.5), 6))
+    plt.bar(x - width/2, em_vals, width, label='Exact Match', color='#3498db')
+    plt.bar(x + width/2, ex_vals, width, label='Execution', color='#2ecc71')
+
+    plt.ylabel('Accuracy (%)', fontweight='bold')
+    plt.title('Model Comparison: Exact Match vs Execution Accuracy', fontweight='bold', fontsize=14)
+    plt.xticks(x, labels, rotation=45, ha="right")
+    plt.legend()
+    plt.ylim(0, max(max(em_vals, default=0), max(ex_vals, default=0)) + 15)
+    plt.grid(axis='y', linestyle='--', alpha=0.7)
+
+    # Attach labels to bars
+    for i in range(len(labels)):
+        plt.text(x[i] - width/2, em_vals[i] + 1, f"{em_vals[i]:.1f}%", ha='center', fontsize=9)
+        plt.text(x[i] + width/2, ex_vals[i] + 1, f"{ex_vals[i]:.1f}%", ha='center', fontsize=9)
+
+    plt.tight_layout()
+    plot_path = plot_dir / "accuracy_comparison.png"
+    plt.savefig(plot_path, dpi=300)
+    print(f"📈 Updated comparison plot saved to: {plot_path}")
+
+if __name__ == "__main__":
+    main()

src/evaluate_model_codet5.py

@@ -0,0 +1,392 @@
+# from __future__ import annotations
+
+# import json
+# import subprocess
+# import sys
+# import argparse
+# import sqlite3
+# import random
+# from pathlib import Path
+
+# import torch
+# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+# from peft import PeftModel
+
+# from prompting import encode_prompt
+
+
+# def _parse_exec_accuracy(stdout: str) -> float | None:
+#     for line in stdout.splitlines():
+#         if line.strip().startswith("execution"):
+#             try:
+#                 return float(line.split()[-1])
+#             except:
+#                 return None
+#     return None
+
+
+# def main():
+
+#     # ---------------- ARGUMENTS ----------------
+#     parser = argparse.ArgumentParser()
+#     parser.add_argument("--adapter", type=str, default="checkpoints/sft_adapter_codet5")
+#     parser.add_argument("--num_samples", type=int, default=1000)
+#     parser.add_argument("--shuffle_dev", action="store_true")
+#     parser.add_argument("--shuffle_seed", type=int, default=42)
+#     parser.add_argument("--accuracy_log", type=str, default="")
+#     args = parser.parse_args()
+
+#     project_root = Path(__file__).resolve().parents[1]
+#     adapter_dir = project_root / args.adapter
+
+#     db_root = project_root / "data" / "database"
+#     table_json = project_root / "data" / "tables.json"
+#     dev_json = project_root / "data" / "dev.json"
+#     gold_sql = project_root / "data" / "dev_gold.sql"
+#     pred_path = project_root / "predictions.txt"
+
+#     if not adapter_dir.exists():
+#         raise FileNotFoundError(f"Missing adapter dir: {adapter_dir}")
+
+#     # ---------------- DEVICE ----------------
+#     device = "mps" if torch.backends.mps.is_available() else (
+#         "cuda" if torch.cuda.is_available() else "cpu"
+#     )
+#     print("Using device:", device)
+
+#     # ---------------- LOAD MODEL ----------------
+#     BASE_MODEL = "Salesforce/codet5-base"
+
+#     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+
+#     if tokenizer.pad_token is None:
+#         tokenizer.pad_token = tokenizer.eos_token
+
+#     base = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL).to(device)
+#     model = PeftModel.from_pretrained(base, str(adapter_dir)).to(device)
+
+#     model = model.merge_and_unload()
+#     model.eval()
+
+#     # ---------------- LOAD DATA ----------------
+#     with dev_json.open() as f:
+#         dev = json.load(f)
+
+#     if args.shuffle_dev:
+#         rng = random.Random(args.shuffle_seed)
+#         rng.shuffle(dev)
+
+#     dev = dev[: args.num_samples]
+
+#     # ---------------- GENERATION CONFIG ----------------
+#     gen_kwargs = dict(
+#         max_new_tokens=160,
+#         num_beams=4,
+#         do_sample=False,
+#         early_stopping=True,
+#         pad_token_id=tokenizer.pad_token_id,
+#         eos_token_id=tokenizer.eos_token_id,
+#     )
+
+#     print("Generating predictions...\n")
+
+#     correct = 0
+#     total = len(dev)
+#     accuracy_log_fh = None
+
+#     if args.accuracy_log:
+#         accuracy_log_path = (project_root / args.accuracy_log).resolve()
+#         accuracy_log_path.parent.mkdir(parents=True, exist_ok=True)
+#         accuracy_log_fh = accuracy_log_path.open("w")
+#         print(f"Writing running accuracy log to: {accuracy_log_path}")
+
+#     with pred_path.open("w") as out_f, torch.no_grad():
+
+#         for i, ex in enumerate(dev, start=1):
+
+#             db_id = ex["db_id"]
+#             question = ex["question"]
+#             gold_query = ex["query"]
+
+#             input_ids = encode_prompt(
+#                 tokenizer,
+#                 question,
+#                 db_id,
+#                 device=device,
+#                 max_input_tokens=512,
+#             )
+
+#             input_ids = input_ids.unsqueeze(0).to(device)
+#             attention_mask = (input_ids != tokenizer.pad_token_id).long().to(device)
+
+#             outputs = model.generate(
+#                 input_ids=input_ids,
+#                 attention_mask=attention_mask,
+#                 **gen_kwargs
+#             )
+
+#             pred_sql = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+#             out_f.write(f"{pred_sql}\t{db_id}\n")
+
+#             # ---------------- LIVE EXECUTION CHECK ----------------
+#             try:
+#                 db_path = db_root / db_id / f"{db_id}.sqlite"
+
+#                 conn = sqlite3.connect(db_path)
+#                 cursor = conn.cursor()
+
+#                 cursor.execute(pred_sql)
+#                 pred_rows = cursor.fetchall()
+
+#                 cursor.execute(gold_query)
+#                 gold_rows = cursor.fetchall()
+
+#                 conn.close()
+
+#                 if sorted(pred_rows) == sorted(gold_rows):
+#                     correct += 1
+
+#             except Exception:
+#                 pass  # execution failed
+
+#             # 🔥 PRINT EVERY 10
+#             if i % 10 == 0 or i == total:
+#                 current_acc = correct / i
+#                 line = f"{i}/{total} | Acc: {current_acc:.3f}"
+#                 print(line)
+#                 if accuracy_log_fh is not None:
+#                     accuracy_log_fh.write(line + "\n")
+
+#     if accuracy_log_fh is not None:
+#         accuracy_log_fh.close()
+
+#     print("\nGeneration finished.\n")
+
+#     # ---------------- OFFICIAL SPIDER EVAL ----------------
+#     eval_script = project_root / "spider_eval" / "evaluation.py"
+
+#     cmd = [
+#         sys.executable,
+#         str(eval_script),
+#         "--gold", str(gold_sql),
+#         "--pred", str(pred_path),
+#         "--etype", "exec",
+#         "--db", str(db_root),
+#         "--table", str(table_json),
+#     ]
+
+#     print("Running Spider evaluation...")
+#     proc = subprocess.run(cmd, capture_output=True, text=True)
+
+#     print(proc.stdout)
+
+#     exec_acc = _parse_exec_accuracy(proc.stdout)
+#     if exec_acc is not None:
+#         print(f"\n🎯 Official Execution Accuracy: {exec_acc*100:.2f}%")
+#     else:
+#         print("Could not parse accuracy.")
+
+
+# if __name__ == "__main__":
+#     main()
+
+import json
+import subprocess
+import sys
+import argparse
+import random
+import sqlite3
+import time
+import re
+from pathlib import Path
+
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+
+# Assuming you have a prompting.py that has encode_prompt
+from prompting import encode_prompt
+
+# -------------------------------
+# LIVE CHECK HELPERS
+# -------------------------------
+def normalize_sql(sql):
+    """Basic normalization for the live progress bar."""
+    sql = sql.replace('"', "'")
+    sql = re.sub(r"\s+", " ", sql)
+    return sql.strip().lower().rstrip(";")
+
+def check_execution(pred_sql, gold_sql, db_path):
+    """Basic execution check for the live progress bar."""
+    try:
+        conn = sqlite3.connect(db_path)
+        conn.text_factory = lambda b: b.decode(errors='ignore')
+        
+        # 2-second timeout so the live tracker doesn't freeze forever
+        start_time = time.monotonic()
+        def timeout_handler():
+            return 1 if (time.monotonic() - start_time) > 2.0 else 0
+        conn.set_progress_handler(timeout_handler, 10000)
+
+        cursor = conn.cursor()
+        cursor.execute(pred_sql)
+        pred_res = cursor.fetchall()
+        
+        cursor.execute(gold_sql)
+        gold_res = cursor.fetchall()
+        conn.close()
+        
+        # Simple sorted check for the live tracker
+        return sorted(pred_res) == sorted(gold_res)
+    except Exception:
+        return False
+
+# -------------------------------
+# SPIDER PARSER
+# -------------------------------
+def _parse_spider_accuracy(stdout: str, metric_type: str) -> float | None:
+    for line in stdout.splitlines():
+        if metric_type == "exec" and line.strip().startswith("execution"):
+            try: return float(line.split()[-1])
+            except: pass
+        elif metric_type == "match" and line.strip().startswith("exact"):
+            try: return float(line.split()[-1])
+            except: pass
+    return None
+
+# -------------------------------
+# MAIN
+# -------------------------------
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--adapter", type=str, required=True, help="Path to your SFT or RLHF checkpoint")
+    parser.add_argument("--num_samples", type=int, default=1034, help="Number of samples to evaluate")
+    parser.add_argument("--shuffle_dev", action="store_true")
+    parser.add_argument("--shuffle_seed", type=int, default=42)
+    args = parser.parse_args()
+
+    project_root = Path(__file__).resolve().parents[1]
+    adapter_dir = project_root / args.adapter
+
+    db_root = project_root / "data" / "database"
+    table_json = project_root / "data" / "tables.json"
+    dev_json = project_root / "data" / "dev.json"
+    
+    pred_path = project_root / "temp_predictions.txt"
+    temp_gold_path = project_root / "temp_gold.sql"
+
+    if not adapter_dir.exists():
+        raise FileNotFoundError(f"Missing adapter dir: {adapter_dir}")
+
+    device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    BASE_MODEL = "Salesforce/codet5-base"
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    print(f"Loading Model: {args.adapter}...")
+    base = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL).to(device)
+    model = PeftModel.from_pretrained(base, str(adapter_dir)).to(device)
+    model = model.merge_and_unload()
+    model.eval()
+
+    with dev_json.open() as f:
+        dev = json.load(f)
+
+    if args.shuffle_dev:
+        rng = random.Random(args.shuffle_seed)
+        rng.shuffle(dev)
+
+    dev = dev[: args.num_samples]
+    total = len(dev)
+
+    gen_kwargs = dict(
+        max_new_tokens=160,
+        num_beams=4,
+        do_sample=False,
+        early_stopping=True,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+
+    print(f"\n🚀 Generating and live-tracking {total} samples...\n")
+
+    em_correct = 0
+    ex_correct = 0
+
+    with pred_path.open("w") as out_pred, temp_gold_path.open("w") as out_gold, torch.no_grad():
+        for i, ex in enumerate(dev, start=1):
+            db_id = ex["db_id"]
+            question = ex["question"]
+            gold_query = ex["query"]
+            db_path = db_root / db_id / f"{db_id}.sqlite"
+
+            # Generate
+            input_ids = encode_prompt(tokenizer, question, db_id, device=device, max_input_tokens=512)
+            input_ids = input_ids.unsqueeze(0).to(device)
+            attention_mask = (input_ids != tokenizer.pad_token_id).long().to(device)
+
+            outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask, **gen_kwargs)
+            pred_sql = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+            
+            # Write to files for official spider eval later
+            out_pred.write(f"{pred_sql}\n")
+            out_gold.write(f"{gold_query}\t{db_id}\n")
+
+            # --- LIVE TRACKING CHECKS ---
+            if normalize_sql(pred_sql) == normalize_sql(gold_query):
+                em_correct += 1
+            if check_execution(pred_sql, gold_query, db_path):
+                ex_correct += 1
+
+            # Print progress every 50 loops
+            if i % 50 == 0 or i == total:
+                print(f"Progress: {i}/{total} | Current EM: {(em_correct/i)*100:.2f}% | Current EX: {(ex_correct/i)*100:.2f}%")
+
+    print("\nGeneration finished. Running Official Spider Evaluations for final numbers...\n")
+
+    eval_script = project_root / "spider_eval" / "evaluation.py"
+
+    # 1. RUN EXACT MATCH EVAL
+    cmd_match = [
+        sys.executable, str(eval_script),
+        "--gold", str(temp_gold_path),
+        "--pred", str(pred_path),
+        "--etype", "match",
+        "--db", str(db_root),
+        "--table", str(table_json),
+    ]
+    proc_match = subprocess.run(cmd_match, capture_output=True, text=True)
+    exact_acc = _parse_spider_accuracy(proc_match.stdout, "match")
+
+    # 2. RUN EXECUTION EVAL
+    cmd_exec = [
+        sys.executable, str(eval_script),
+        "--gold", str(temp_gold_path),
+        "--pred", str(pred_path),
+        "--etype", "exec",
+        "--db", str(db_root),
+        "--table", str(table_json),
+    ]
+    proc_exec = subprocess.run(cmd_exec, capture_output=True, text=True)
+    exec_acc = _parse_spider_accuracy(proc_exec.stdout, "exec")
+
+    print("==========================================")
+    print(f"🎯 OFFICIAL SPIDER RESULTS FOR: {args.adapter}")
+    print("==========================================")
+    
+    if exact_acc is not None:
+        print(f"Exact Set Match Accuracy  : {exact_acc*100:.2f}%")
+    else:
+        print("Exact Set Match Accuracy  : Could not parse output")
+        
+    if exec_acc is not None:
+        print(f"Execution Accuracy        : {exec_acc*100:.2f}%")
+    else:
+        print("Execution Accuracy        : Could not parse output")
+    print("==========================================\n")
+
+if __name__ == "__main__":
+    main()

src/evaluate_model_t5_small_sft.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+import argparse
+import re
+import sqlite3
+from pathlib import Path
+
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+from prompting import encode_prompt
+
+
+# ---------------- PARSE ACC ----------------
+def _parse_exec_accuracy(stdout: str) -> float | None:
+    for line in stdout.splitlines():
+        if line.strip().startswith("execution"):
+            try:
+                return float(line.split()[-1])
+            except:
+                return None
+    return None
+
+
+# ---------------- CLEAN SQL ----------------
+def clean_prediction(pred_sql: str) -> str:
+    pred_sql = pred_sql.strip()
+
+    if "SQL:" in pred_sql:
+        pred_sql = pred_sql.split("SQL:")[-1]
+
+    pred_sql = pred_sql.replace('"', "'")
+    pred_sql = re.sub(r"\s+", " ", pred_sql).strip()
+
+    if not pred_sql.endswith(";"):
+        pred_sql += ";"
+
+    return pred_sql
+
+
+def main():
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--adapter", type=str, default="checkpoints/sft_t5")
+    parser.add_argument("--num_samples", type=int, default=1000)
+    args = parser.parse_args()
+
+    project_root = Path(__file__).resolve().parents[1]
+    adapter_dir = project_root / args.adapter
+
+    db_root = project_root / "data/database"
+    table_json = project_root / "data/tables.json"
+    dev_json = project_root / "data/dev.json"
+    gold_sql = project_root / "data/dev_gold.sql"
+    pred_path = project_root / "pred.sql"
+
+    if not adapter_dir.exists():
+        raise FileNotFoundError(f"Missing adapter dir: {adapter_dir}")
+
+    # ---------------- DEVICE ----------------
+    device = "mps" if torch.backends.mps.is_available() else (
+        "cuda" if torch.cuda.is_available() else "cpu"
+    )
+    print("Using device:", device)
+
+    # ---------------- LOAD MODEL ----------------
+    BASE_MODEL = "t5-small"
+
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+    base = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL).to(device)
+
+    model = PeftModel.from_pretrained(base, str(adapter_dir)).to(device)
+    model = model.merge_and_unload()
+    model.eval()
+
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # ---------------- LOAD DATA ----------------
+    with dev_json.open() as f:
+        dev = json.load(f)[: args.num_samples]
+
+    print("Generating predictions...\n")
+
+    correct = 0
+    total = len(dev)
+
+    # ---------------- GENERATE + LIVE EXEC ----------------
+    with pred_path.open("w") as out_f, torch.no_grad():
+
+        for i, ex in enumerate(dev, start=1):
+
+            db_id = ex["db_id"]
+            question = ex["question"]
+            gold_query = ex["query"]
+
+            prompt_ids = encode_prompt(
+                tokenizer,
+                question,
+                db_id,
+                device=device,
+                max_input_tokens=512,
+            )
+
+            input_ids = prompt_ids.unsqueeze(0).to(device)
+            attention_mask = (input_ids != tokenizer.pad_token_id).long().to(device)
+
+            outputs = model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_new_tokens=160,
+                num_beams=4,
+                do_sample=False,
+                early_stopping=True,
+            )
+
+            pred_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            pred_sql = clean_prediction(pred_sql)
+
+            out_f.write(pred_sql + "\n")
+
+            # -------- LIVE EXECUTION CHECK --------
+            try:
+                db_path = db_root / db_id / f"{db_id}.sqlite"
+
+                conn = sqlite3.connect(db_path)
+                cursor = conn.cursor()
+
+                cursor.execute(pred_sql)
+                pred_rows = cursor.fetchall()
+
+                cursor.execute(gold_query)
+                gold_rows = cursor.fetchall()
+
+                conn.close()
+
+                if sorted(pred_rows) == sorted(gold_rows):
+                    correct += 1
+
+            except Exception:
+                pass  # execution failed
+
+            # 🔥 PRINT EVERY 10
+            if i % 10 == 0 or i == total:
+                current_acc = correct / i
+                print(f"{i}/{total} | Acc: {current_acc:.3f}")
+
+    print("\nGeneration finished.\n")
+
+    # ---------------- SPIDER EVAL ----------------
+    eval_script = project_root / "spider_eval/evaluation.py"
+
+    cmd = [
+        sys.executable,
+        str(eval_script),
+        "--gold", str(gold_sql),
+        "--pred", str(pred_path),
+        "--etype", "exec",
+        "--db", str(db_root),
+        "--table", str(table_json),
+    ]
+
+    print("Running Spider evaluation...")
+    proc = subprocess.run(cmd, capture_output=True, text=True)
+
+    print(proc.stdout)
+
+    exec_acc = _parse_exec_accuracy(proc.stdout)
+    if exec_acc is not None:
+        print(f"\n🎯 Official Execution Accuracy: {exec_acc*100:.2f}%")
+    else:
+        print("Could not parse accuracy.")
+
+
+if __name__ == "__main__":
+    main()

src/evaluate_rl_bart.py

@@ -0,0 +1,138 @@
+
+import json
+import sqlite3
+import argparse
+import time
+from pathlib import Path
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+
+# ---------------- PROMPT (IDENTICAL TO TRAINING) ----------------
+def build_prompt(question, schema):
+    return f"""
+Database Schema:
+{schema}
+
+Translate English to SQL:
+{question}
+SQL:
+"""
+
+# ---------------- LOAD SCHEMA ----------------
+def load_schema(db_path):
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    tables = cursor.execute(
+        "SELECT name FROM sqlite_master WHERE type='table';"
+    ).fetchall()
+
+    schema = ""
+    for (table,) in tables:
+        cols = cursor.execute(f"PRAGMA table_info({table});").fetchall()
+        col_names = [c[1] for c in cols]
+        schema += f"{table}({', '.join(col_names)})\n"
+
+    conn.close()
+    return schema
+
+
+# ---------------- EXECUTION CHECK WITH TIMEOUT ----------------
+def execution_match(pred_sql, gold_sql, db_path):
+    try:
+        conn = sqlite3.connect(db_path)
+        
+        # --- 5-SECOND TIMEOUT SO EVALUATION DOESN'T FREEZE ---
+        start_time = time.monotonic()
+        def timeout_handler():
+            return 1 if (time.monotonic() - start_time) > 5.0 else 0
+        conn.set_progress_handler(timeout_handler, 10000)
+
+        cur = conn.cursor()
+
+        cur.execute(pred_sql)
+        pred = cur.fetchall()
+
+        cur.execute(gold_sql)
+        gold = cur.fetchall()
+
+        conn.close()
+        return pred == gold
+
+    except Exception:
+        return False
+
+
+# ---------------- MAIN ----------------
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--adapter", type=str, required=True)
+    parser.add_argument("--num_samples", type=int, default=1034)
+    args = parser.parse_args()
+
+    project_root = Path(__file__).resolve().parents[1]
+
+    dev_json = project_root / "data" / "dev.json"
+    db_root = project_root / "data" / "database"
+
+    # 🎯 Added CUDA support for Nvidia GPUs
+    device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
+
+    # load model
+    base_model = "facebook/bart-base"
+    print(f"Loading Base: {base_model}")
+    print(f"Loading Adapter: {args.adapter}")
+     
+    tokenizer = AutoTokenizer.from_pretrained(args.adapter)
+    base = AutoModelForSeq2SeqLM.from_pretrained(base_model).to(device)
+    model = PeftModel.from_pretrained(base, args.adapter).to(device)
+    model = model.merge_and_unload()
+
+    with open(dev_json) as f:
+        dev = json.load(f)[: args.num_samples]
+
+    correct = 0
+
+    print(f"Evaluating {len(dev)} examples...\n")
+
+    for i, ex in enumerate(dev, 1):
+        question = ex["question"]
+        db_id = ex["db_id"]
+        gold_sql = ex["query"]
+
+        db_path = db_root / db_id / f"{db_id}.sqlite"
+        schema = load_schema(db_path)
+
+        prompt = build_prompt(question, schema)
+
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=80,
+                do_sample=False,
+                num_beams=4,
+            )
+
+        pred_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        if "SQL:" in pred_sql:
+            pred_sql = pred_sql.split("SQL:")[-1].strip()
+
+        match = execution_match(pred_sql, gold_sql, db_path)
+
+        if match:
+            correct += 1
+
+        if i % 10 == 0:
+            print(f"{i}/{len(dev)} | Acc: {correct/i:.3f}")
+
+    print("\n=============================")
+    print(f"FINAL EXECUTION ACCURACY: {correct/len(dev)*100:.2f}%")
+    print("=============================")
+
+
+if __name__ == "__main__":
+    main()

src/evaluate_sft_bart.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+import argparse
+import re
+import sqlite3
+from pathlib import Path
+
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+from prompting import encode_prompt
+
+
+# ---------------- SQL CLEAN ----------------
+def extract_sql(text: str) -> str:
+    text = text.strip()
+
+    if "SQL:" in text:
+        text = text.split("SQL:")[-1]
+
+    match = re.search(r"(SELECT .*?)(?:$)", text, re.IGNORECASE | re.DOTALL)
+    if match:
+        text = match.group(1)
+
+    text = text.replace('"', "'")
+    text = re.sub(r"\s+", " ", text).strip()
+
+    if not text.endswith(";"):
+        text += ";"
+
+    return text
+
+
+# ---------------- ROBUST ACC PARSER ----------------
+def parse_exec_accuracy(stdout: str):
+    for line in stdout.splitlines():
+        if "execution" in line.lower():
+            numbers = re.findall(r"\d+\.\d+", line)
+            if numbers:
+                return float(numbers[-1])
+    return None
+
+
+def main():
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--adapter", type=str, default="checkpoints/sft_best_bart_2")
+    parser.add_argument("--num_samples", type=int, default=1000)
+    args = parser.parse_args()
+
+    project_root = Path(__file__).resolve().parents[1]
+    adapter_dir = project_root / args.adapter
+
+    if not adapter_dir.exists():
+        raise FileNotFoundError(f"Adapter not found: {adapter_dir}")
+
+    db_root = project_root / "data/database"
+    table_json = project_root / "data/tables.json"
+    dev_json = project_root / "data/dev.json"
+    gold_sql_file = project_root / "data/dev_gold.sql"
+    pred_sql_file = project_root / "pred.sql"
+
+    device = "mps" if torch.backends.mps.is_available() else (
+        "cuda" if torch.cuda.is_available() else "cpu"
+    )
+    print("Using device:", device)
+
+    # -------- LOAD MODEL --------
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(adapter_dir)
+
+    BASE_MODEL = "facebook/bart-base"
+    print(f"Loading base model {BASE_MODEL}...")
+    base_model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL).to(device)
+
+    print("Loading LoRA adapter...")
+    model = PeftModel.from_pretrained(base_model, adapter_dir).to(device)
+    model = model.merge_and_unload()
+    model.eval()
+
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # -------- LOAD DATA --------
+    with open(dev_json) as f:
+        dev = json.load(f)[: args.num_samples]
+
+    print("Generating SQL predictions...\n")
+
+    correct = 0
+    total = len(dev)
+
+    with open(pred_sql_file, "w") as f, torch.no_grad():
+
+        for i, ex in enumerate(dev, 1):
+
+            question = ex["question"]
+            db_id = ex["db_id"]
+            gold_query = ex["query"]
+
+            prompt_ids = encode_prompt(
+                tokenizer,
+                question,
+                db_id,
+                device=device,
+                max_input_tokens=512,
+            )
+
+            input_ids = prompt_ids.unsqueeze(0).to(device)
+            attention_mask = (input_ids != tokenizer.pad_token_id).long().to(device)
+
+            outputs = model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_new_tokens=160,
+                num_beams=4,
+                do_sample=False,
+            )
+
+            pred = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            pred_sql = extract_sql(pred)
+
+            f.write(f"{pred_sql}\t{db_id}\n")
+
+            # -------- LIVE EXECUTION CHECK --------
+            try:
+                db_path = db_root / db_id / f"{db_id}.sqlite"
+
+                conn = sqlite3.connect(db_path)
+                cursor = conn.cursor()
+
+                cursor.execute(pred_sql)
+                pred_rows = cursor.fetchall()
+
+                cursor.execute(gold_query)
+                gold_rows = cursor.fetchall()
+
+                conn.close()
+
+                # order insensitive comparison
+                if sorted(pred_rows) == sorted(gold_rows):
+                    correct += 1
+
+            except Exception:
+                pass  # execution failed
+
+            if i % 10 == 0 or i == total:
+                current_acc = correct / i
+                print(f"{i}/{total} | Acc: {current_acc:.3f}")
+
+    print("\nGeneration finished.\n")
+
+    # -------- RUN OFFICIAL SPIDER EVAL --------
+    eval_script = project_root / "spider_eval/evaluation.py"
+    if (project_root / "spider_eval/evaluation_bart.py").exists():
+        eval_script = project_root / "spider_eval/evaluation_bart.py"
+
+    cmd = [
+        sys.executable,
+        str(eval_script),
+        "--gold", str(gold_sql_file),
+        "--pred", str(pred_sql_file),
+        "--etype", "exec",
+        "--db", str(db_root),
+        "--table", str(table_json),
+    ]
+
+    print(f"\nRunning Spider evaluation using {eval_script.name}...")
+    proc = subprocess.run(cmd, capture_output=True, text=True, errors="ignore")
+
+    if proc.returncode != 0:
+        print("\nSpider evaluation crashed.")
+        print(proc.stderr)
+        return
+
+    print("\n--- Spider Eval Output ---")
+    print("\n".join(proc.stdout.splitlines()[-20:]))
+
+    acc = parse_exec_accuracy(proc.stdout)
+    if acc is not None:
+        print(f"\n🎯 Official Execution Accuracy: {acc*100:.2f}%")
+    else:
+        print("\nCould not parse official accuracy.")
+
+
+if __name__ == "__main__":
+    main()

src/execution_reward.py

@@ -0,0 +1,409 @@
+from __future__ import annotations
+
+import os
+import re
+import sqlite3
+import time
+from dataclasses import dataclass
+from typing import List, Optional, Sequence, Set, Tuple, Union
+
+try:
+    import sqlparse
+    from sqlparse.sql import Function, Identifier, IdentifierList, Statement, Token, Where
+    from sqlparse.tokens import DML, Keyword, Name, Number, Punctuation, String, Whitespace
+except Exception:  # pragma: no cover
+    sqlparse = None  # type: ignore[assignment]
+    Statement = object  # type: ignore[misc,assignment]
+    Token = object  # type: ignore[misc,assignment]
+
+
+def _normalize_sql(sql: str) -> str:
+    if not isinstance(sql, str):
+        return ""
+    s = sql.strip()
+    if s.startswith("```"):
+        # Strip markdown fences if present.
+        s = re.sub(r"^```[a-zA-Z0-9_+-]*\n?", "", s).strip()
+        s = re.sub(r"\n?```$", "", s).strip()
+    if s.lower().startswith("sql:"):
+        s = s[4:].strip()
+    # Keep only the first statement to avoid accidental multi-statement execution.
+    if ";" in s:
+        s = s.split(";", 1)[0].strip()
+    return s
+
+
+def _connect_readonly(db_path: str) -> sqlite3.Connection:
+    # Read-only prevents any accidental mutation during reward computation.
+    # Note: requires SQLite URI support (built-in).
+    uri = f"file:{os.path.abspath(db_path)}?mode=ro"
+    conn = sqlite3.connect(uri, uri=True, check_same_thread=False)
+    conn.execute("PRAGMA query_only = ON;")
+    conn.execute("PRAGMA foreign_keys = ON;")
+    return conn
+
+
+def _with_timeout(conn: sqlite3.Connection, timeout_s: float = 1.0) -> None:
+    start = time.monotonic()
+
+    def _handler() -> int:
+        return 1 if (time.monotonic() - start) > timeout_s else 0
+
+    # Call handler every N VM opcodes.
+    conn.set_progress_handler(_handler, 10_000)
+
+
+def _list_tables(conn: sqlite3.Connection) -> List[str]:
+    try:
+        cur = conn.execute(
+            "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';"
+        )
+        return [r[0] for r in cur.fetchall() if r and isinstance(r[0], str)]
+    except sqlite3.Error:
+        return []
+
+
+def _contains_table_name(sql: str, table_names: Sequence[str]) -> bool:
+    s = sql.lower()
+    for t in table_names:
+        tl = t.lower()
+        if not tl:
+            continue
+        if re.search(rf"\b{re.escape(tl)}\b", s):
+            return True
+    return False
+
+
+def _explain_query_plan(conn: sqlite3.Connection, sql: str) -> bool:
+    try:
+        _with_timeout(conn, timeout_s=1.0)
+        conn.execute(f"EXPLAIN QUERY PLAN {sql}")
+        return True
+    except sqlite3.Error:
+        return False
+
+
+def _execute(conn: sqlite3.Connection, sql: str, max_rows: int = 1000) -> Tuple[bool, List[Tuple], Optional[str]]:
+    try:
+        _with_timeout(conn, timeout_s=1.0)
+        cur = conn.execute(sql)
+        rows = cur.fetchmany(max_rows)
+        # Normalize to plain tuples for deterministic comparison.
+        norm_rows = [tuple(r) for r in rows]
+        return True, norm_rows, None
+    except sqlite3.Error as e:
+        return False, [], str(e)
+
+
+_SQL_KEYWORDS_TO_IGNORE = {
+    "select",
+    "from",
+    "where",
+    "join",
+    "inner",
+    "left",
+    "right",
+    "full",
+    "outer",
+    "on",
+    "group",
+    "by",
+    "order",
+    "limit",
+    "having",
+    "distinct",
+    "union",
+    "intersect",
+    "except",
+    "as",
+    "and",
+    "or",
+    "not",
+    "in",
+    "is",
+    "null",
+    "like",
+    "between",
+    "case",
+    "when",
+    "then",
+    "else",
+    "end",
+    "asc",
+    "desc",
+}
+
+_SQL_FUNCTIONS_TO_IGNORE = {
+    "count",
+    "avg",
+    "min",
+    "max",
+    "sum",
+    "lower",
+    "upper",
+    "substr",
+    "coalesce",
+    "round",
+    "date",
+    "datetime",
+    "strftime",
+}
+
+
+def extract_tables(sql: str) -> Set[str]:
+    """
+    Best-effort table extraction from SQL using sqlparse.
+    Returns lowercase table names (unqualified).
+    """
+    sql = _normalize_sql(sql)
+    if not sql:
+        return set()
+    if sqlparse is None:
+        # Fallback: naive regex for FROM/JOIN.
+        found = set()
+        for m in re.finditer(r"\b(from|join)\s+([a-zA-Z_][\w$]*)", sql, flags=re.I):
+            found.add(m.group(2).lower())
+        return found
+
+    try:
+        statements = sqlparse.parse(sql)
+    except Exception:
+        return set()
+
+    tables: Set[str] = set()
+
+    def _add_identifier_as_table(ident: Identifier) -> None:
+        # Prefer real name over alias; strip any schema prefix.
+        name = ident.get_real_name() or ident.get_name()
+        if not name:
+            return
+        tables.add(name.lower())
+
+    for st in statements:
+        if not isinstance(st, Statement):
+            continue
+        seen_from = False
+        for tok in st.flatten():
+            if tok.ttype in Whitespace:
+                continue
+            if tok.ttype is Keyword and tok.value.upper() in {"FROM", "JOIN", "INNER JOIN", "LEFT JOIN", "RIGHT JOIN", "FULL JOIN"}:
+                seen_from = True
+                continue
+            if not seen_from:
+                continue
+
+            if isinstance(tok, Identifier):
+                _add_identifier_as_table(tok)
+                seen_from = False
+            elif tok.ttype is Name:
+                tables.add(tok.value.lower())
+                seen_from = False
+            elif tok.ttype is Keyword and tok.value.upper() in {"WHERE", "GROUP", "ORDER", "HAVING", "LIMIT"}:
+                seen_from = False
+
+    return tables
+
+
+def extract_columns(sql: str) -> Set[str]:
+    """
+    Best-effort column extraction from SQL using sqlparse.
+    Returns lowercase column names (unqualified).
+    """
+    sql = _normalize_sql(sql)
+    if not sql:
+        return set()
+    if sqlparse is None:
+        # Fallback: naive dotted identifiers and bare names after SELECT/WHERE/etc.
+        cols = set()
+        for m in re.finditer(r"\b([a-zA-Z_][\w$]*)\b", sql):
+            w = m.group(1).lower()
+            if w in _SQL_KEYWORDS_TO_IGNORE or w in _SQL_FUNCTIONS_TO_IGNORE:
+                continue
+            cols.add(w)
+        return cols
+
+    try:
+        statements = sqlparse.parse(sql)
+    except Exception:
+        return set()
+
+    cols: Set[str] = set()
+
+    def _maybe_add_col(name: Optional[str]) -> None:
+        if not name:
+            return
+        n = name.strip().strip('"').strip("'").lower()
+        if not n or n == "*":
+            return
+        if n in _SQL_KEYWORDS_TO_IGNORE or n in _SQL_FUNCTIONS_TO_IGNORE:
+            return
+        cols.add(n)
+
+    def _handle_identifier(ident: Identifier) -> None:
+        # If qualified (t.col), keep only col for overlap/hallucination checks.
+        _maybe_add_col(ident.get_real_name() or ident.get_name())
+
+    for st in statements:
+        if not isinstance(st, Statement):
+            continue
+        for tok in st.flatten():
+            # Skip whitespace/punctuation/string literals/numbers.
+            if getattr(tok, "ttype", None) in (Whitespace, Punctuation, String, Number):
+                continue
+
+            if isinstance(tok, Function):
+                fname = tok.get_name()
+                if fname:
+                    # Don't treat function name as a column.
+                    pass
+                continue
+
+            if isinstance(tok, IdentifierList):
+                for ident in tok.get_identifiers():
+                    if isinstance(ident, Identifier):
+                        _handle_identifier(ident)
+                continue
+
+            if isinstance(tok, Identifier):
+                _handle_identifier(tok)
+                continue
+
+            if getattr(tok, "ttype", None) is Name:
+                _maybe_add_col(tok.value)
+
+    return cols
+
+
+def _get_db_tables_and_columns(conn: sqlite3.Connection) -> Tuple[Set[str], Set[str]]:
+    """
+    Return (tables, columns) sets from SQLite schema; all lowercased.
+    Columns are returned as a global set (unqualified).
+    """
+    tables = set()
+    columns = set()
+    for t in _list_tables(conn):
+        tl = t.lower()
+        if not tl:
+            continue
+        tables.add(tl)
+        try:
+            cur = conn.execute(f'PRAGMA table_info("{t}")')
+            for row in cur.fetchall():
+                if row and isinstance(row[1], str):
+                    columns.add(row[1].lower())
+        except sqlite3.Error:
+            continue
+    return tables, columns
+
+
+def _safe_results_equal(a: List[Tuple], b: List[Tuple]) -> bool:
+    # Deterministic comparison: compare exact row tuples in order.
+    return a == b
+
+
+@dataclass
+class RewardDebugStats:
+    total: int = 0
+    parsed_ok: int = 0
+    table_match: int = 0
+    column_match: int = 0
+    executed_ok: int = 0
+    exact_match: int = 0
+
+
+_DEBUG = RewardDebugStats()
+
+
+def reset_debug_metrics() -> None:
+    global _DEBUG
+    _DEBUG = RewardDebugStats()
+
+
+def get_debug_metrics() -> dict:
+    denom = max(_DEBUG.total, 1)
+    return {
+        "valid_sql_rate": _DEBUG.parsed_ok / denom,
+        "table_match_rate": _DEBUG.table_match / denom,
+        "column_match_rate": _DEBUG.column_match / denom,
+        "execution_accuracy": _DEBUG.exact_match / denom,
+    }
+
+EXECUTION_ERROR = "EXECUTION_ERROR"
+
+
+def execute_sql(conn: sqlite3.Connection, sql: str, *, max_rows: int = 1000) -> Union[List[Tuple], str]:
+    """
+    Execute SQL safely.
+
+    If sqlite raises ANY exception, return EXECUTION_ERROR (NOT empty list).
+    """
+    try:
+        _with_timeout(conn, timeout_s=1.0)
+        cur = conn.execute(sql)
+        rows = cur.fetchmany(max_rows)
+        return [tuple(r) for r in rows]
+    except Exception:
+        return EXECUTION_ERROR
+
+
+def _sqlparse_valid_select(sql: str) -> bool:
+    """
+    Parse validation using sqlparse:
+      - parse() non-empty
+      - contains a SELECT statement
+    """
+    if sqlparse is None:
+        return False
+    try:
+        stmts = sqlparse.parse(sql)
+        if not stmts:
+            return False
+        for st in stmts:
+            try:
+                if hasattr(st, "get_type") and st.get_type() == "SELECT":
+                    return True
+            except Exception:
+                continue
+        return False
+    except Exception:
+        return False
+
+def execution_reward(pred_sql: str, db_path: str, gold_sql: str) -> float:
+    try:
+        sql = _normalize_sql(pred_sql)
+        gold = _normalize_sql(gold_sql)
+
+        if not sql or "SELECT" not in sql.upper():
+            return -1.0
+
+        if not _sqlparse_valid_select(sql):
+            return -1.0
+
+        reward = -0.2  # valid SQL baseline
+
+        pred_tables = extract_tables(sql)
+        gold_tables = extract_tables(gold)
+
+        if pred_tables == gold_tables and len(gold_tables) > 0:
+            reward += 0.3
+
+        pred_cols = extract_columns(sql)
+        gold_cols = extract_columns(gold)
+
+        if gold_cols:
+            overlap = len(pred_cols & gold_cols) / len(gold_cols)
+            reward += 0.3 * overlap
+
+        with _connect_readonly(db_path) as conn:
+            pred_res = execute_sql(conn, sql)
+            if pred_res != EXECUTION_ERROR:
+                reward += 0.2
+
+            gold_res = execute_sql(conn, gold)
+            if pred_res != EXECUTION_ERROR and _safe_results_equal(pred_res, gold_res):
+                return 1.0
+
+        return max(-1.0, min(1.0, reward))
+
+    except Exception:
+        return -1.0

src/generate_sql.py

@@ -0,0 +1,68 @@
+import argparse
+import os
+
+import torch
+from transformers import AutoTokenizer
+from transformers import AutoModelForSeq2SeqLM
+from peft import PeftModel
+
+from prompting import encode_prompt
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate SQL from a question + db_id using the RLHF model.")
+    parser.add_argument("--question", type=str, required=True)
+    parser.add_argument("--db_id", type=str, required=True)
+    parser.add_argument("--model_dir", type=str, default=None, help="Defaults to outputs/rlhf_text2sql")
+    parser.add_argument("--use_schema", action="store_true", help="Include schema in the prompt (must match training).")
+    parser.add_argument("--max_schema_chars", type=int, default=1500)
+    parser.add_argument("--max_new_tokens", type=int, default=80)
+    args = parser.parse_args()
+
+    device = "mps" if torch.backends.mps.is_available() else "cpu"
+
+    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    adapter_dir = args.model_dir or os.path.join(project_root, "outputs", "rlhf_text2sql")
+    base_model = os.environ.get("BASE_MODEL", "t5-small")
+    fallback_base_model = os.path.join(project_root, "models", "t5_spider_sft")
+    if not os.path.isdir(base_model) and os.path.isdir(fallback_base_model):
+        base_model = fallback_base_model
+
+    local_only = not os.path.isdir(base_model)
+    tokenizer_source = adapter_dir if os.path.isdir(adapter_dir) else base_model
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_source, local_files_only=not os.path.isdir(tokenizer_source))
+    base = AutoModelForSeq2SeqLM.from_pretrained(base_model, local_files_only=local_only).to(device)
+    model = PeftModel.from_pretrained(base, adapter_dir).to(device)
+    # Merge adapters for faster/stabler generation.
+    model = model.merge_and_unload()
+    model.config.use_cache = False
+
+    if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    input_ids = encode_prompt(
+        tokenizer,
+        args.question,
+        args.db_id,
+        device=device,
+        max_input_tokens=512,
+    )
+
+    gen_kwargs = dict(
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+        num_beams=1,
+        early_stopping=True,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+
+    with torch.no_grad():
+        out = model.generate(input_ids=input_ids.unsqueeze(0), **gen_kwargs)
+
+    sql = tokenizer.decode(out[0], skip_special_tokens=True).strip()
+    print(sql)
+
+
+if __name__ == "__main__":
+    main()

src/human_eval_runner.py

@@ -0,0 +1,152 @@
+import json
+import sqlite3
+from pathlib import Path
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+DB_ROOT = PROJECT_ROOT / "data" / "database"
+
+# Added CUDA fallback for consistency
+DEVICE = "mps" if torch.backends.mps.is_available() else (
+    "cuda" if torch.cuda.is_available() else "cpu"
+)
+
+# ================= LOAD MODEL =================
+def load_model(adapter_path):
+    base_name = "Salesforce/codet5-base"
+
+    # 🐛 FIXED: Convert relative path to absolute path to prevent Hugging Face 404 errors
+    abs_path = (PROJECT_ROOT / adapter_path).resolve()
+    if not abs_path.exists():
+        raise FileNotFoundError(f"Adapter not found at: {abs_path}")
+
+    print(f"\nLoading model from: {abs_path}")
+
+    # 🐛 FIXED: Added fallback in case tokenizer isn't saved in the adapter folder
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(str(abs_path), local_files_only=True)
+    except Exception:
+        print("Adapter tokenizer missing — using base tokenizer")
+        tokenizer = AutoTokenizer.from_pretrained(base_name)
+
+    base = AutoModelForSeq2SeqLM.from_pretrained(base_name).to(DEVICE)
+    model = PeftModel.from_pretrained(base, str(abs_path)).to(DEVICE)
+    model.eval()
+
+    return tokenizer, model
+
+
+# ================= SCHEMA =================
+def load_schema(db_id):
+    db_path = DB_ROOT / db_id / f"{db_id}.sqlite"
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    tables = cursor.execute(
+        "SELECT name FROM sqlite_master WHERE type='table';"
+    ).fetchall()
+
+    schema = ""
+    for (table,) in tables:
+        cols = cursor.execute(f"PRAGMA table_info({table});").fetchall()
+        col_names = [c[1] for c in cols]
+        schema += f"{table}({', '.join(col_names)})\n"
+
+    conn.close()
+    return schema
+
+
+# ================= GENERATE =================
+def generate_sql(tokenizer, model, question, db_id):
+    schema = load_schema(db_id)
+
+    prompt = f"""
+Database Schema:
+{schema}
+
+Translate English to SQL:
+{question}
+SQL:
+"""
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=120,
+            num_beams=4,
+            do_sample=False
+        )
+
+    sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    if "SQL:" in sql:
+        sql = sql.split("SQL:")[-1]
+
+    return sql.strip()
+
+
+# ================= EXECUTE =================
+def try_execute(sql, db_id):
+    db_path = DB_ROOT / db_id / f"{db_id}.sqlite"
+    try:
+        conn = sqlite3.connect(db_path)
+        cur = conn.cursor()
+        cur.execute(sql)
+        cur.fetchall()
+        conn.close()
+        return True
+    except:
+        return False
+
+
+# ================= MAIN =================
+def main():
+    # paths (change if needed)
+    SFT_MODEL = "checkpoints/sft_adapter_codet5" # Ensure this matches your actual SFT folder name!
+    RLHF_MODEL = "checkpoints/best_rlhf_model"
+
+    tokenizer_sft, model_sft = load_model(SFT_MODEL)
+    tokenizer_rl, model_rl = load_model(RLHF_MODEL)
+
+    human_eval_path = PROJECT_ROOT / "data/human_eval.json"
+    with open(human_eval_path) as f:
+        questions = json.load(f)
+
+    sft_success = 0
+    rl_success = 0
+
+    print("\nRunning Human Evaluation...\n")
+
+    for i, q in enumerate(questions, 1):
+        db = q["db_id"]
+        question = q["question"]
+
+        sql_sft = generate_sql(tokenizer_sft, model_sft, question, db)
+        sql_rl = generate_sql(tokenizer_rl, model_rl, question, db)
+
+        ok_sft = try_execute(sql_sft, db)
+        ok_rl = try_execute(sql_rl, db)
+
+        if ok_sft:
+            sft_success += 1
+        if ok_rl:
+            rl_success += 1
+
+        print(f"\nQ{i}: {question}")
+        print(f"SFT : {'OK' if ok_sft else 'FAIL'}")
+        print(f"RLHF: {'OK' if ok_rl else 'FAIL'}")
+
+    print("\n=============================")
+    print("HUMAN EVALUATION RESULT")
+    print("=============================")
+    print(f"SFT  Success: {sft_success}/{len(questions)} = {sft_success/len(questions)*100:.2f}%")
+    print(f"RLHF Success: {rl_success}/{len(questions)} = {rl_success/len(questions)*100:.2f}%")
+    print("=============================\n")
+
+
+if __name__ == "__main__":
+    main()

src/load_lora_model.py

@@ -0,0 +1,30 @@
+import torch
+from transformers import T5ForConditionalGeneration, T5Tokenizer
+from peft import LoraConfig, get_peft_model, TaskType
+
+device = "mps" if torch.backends.mps.is_available() else "cpu"
+
+MODEL_PATH = "../outputs/model"   # your supervised trained model
+
+print("Loading base model...")
+model = T5ForConditionalGeneration.from_pretrained(MODEL_PATH).to(device)
+
+tokenizer = T5Tokenizer.from_pretrained("t5-small")
+
+# ---------------- LoRA CONFIG ----------------
+lora_config = LoraConfig(
+    r=8,                       # rank (small brain attachment)
+    lora_alpha=16,
+    target_modules=["q", "v"], # attention matrices only
+    lora_dropout=0.05,
+    bias="none",
+    task_type=TaskType.SEQ_2_SEQ_LM
+)
+
+print("Attaching LoRA adapters...")
+model = get_peft_model(model, lora_config)
+
+model.print_trainable_parameters()
+
+print("READY ✔ LoRA model loaded")
+

src/make_rl_dataset.py

@@ -0,0 +1,20 @@
+import json
+from datasets import load_dataset
+
+print("Loading Spider dataset...")
+dataset = load_dataset("spider", split="train")
+
+data = []
+
+for ex in dataset:
+    data.append({
+        "question": ex["question"],
+        "query": ex["query"],
+        "db_id": ex["db_id"]   # ⭐ CRITICAL FIELD
+    })
+
+print("Saving JSON...")
+with open("data/train_spider.json", "w") as f:
+    json.dump(data, f, indent=2)
+
+print("Done! File saved at data/train_spider.json")

src/manual_check.py

@@ -0,0 +1,44 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+
+BASE_MODEL = "Salesforce/codet5-base"
+ADAPTER = "checkpoints/sft_adapter"   # change if needed
+
+device = "mps" if torch.backends.mps.is_available() else "cpu"
+
+print("Loading model...")
+
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)
+model = PeftModel.from_pretrained(model, ADAPTER)
+
+model = model.to(device)
+model.eval()
+
+# 5 random Spider style questions
+questions = [
+    "List all employee names",
+    "Find the number of students in each department",
+    "Show the average salary of employees",
+    "Which flights depart from LA?",
+    "Find customers who bought more than 5 items"
+]
+
+for q in questions:
+    prompt = f"Translate to SQL: {q}"
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=128,
+            temperature=0.0,   # deterministic
+        )
+
+    sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    print("\nQUESTION:", q)
+    print("SQL:", sql)
+    print("-"*60)

src/predict.py

@@ -0,0 +1,112 @@
+import torch
+import sqlite3
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+# --------------------------------------------------
+# PATH
+# --------------------------------------------------
+MODEL_PATH = "outputs/model"
+
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+
+print("Loading fine-tuned model...")
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH)
+model.eval()
+
+# --------------------------------------------------
+# CONNECT DATABASE
+# --------------------------------------------------
+print("Connecting to database...")
+# conn = sqlite3.connect("../data/database/department_management/department_management.sqlite")
+conn = sqlite3.connect("data/database/department_management/department_management.sqlite")
+cursor = conn.cursor()
+print("Database connected ✔")
+
+# --------------------------------------------------
+# BUILD PROMPT
+# --------------------------------------------------
+def build_prompt(question):
+    schema = """
+Table department columns = Department_ID, Name, Creation, Ranking, Budget_in_Billions, Num_Employees.
+Table head columns = head_ID, name, born_state, age.
+Table management columns = department_ID, head_ID, temporary_acting.
+"""
+    return f"translate English to SQL: {schema} question: {question}"
+
+# --------------------------------------------------
+# GENERATE SQL
+# --------------------------------------------------
+def generate_sql(question):
+
+    prompt = build_prompt(question)
+
+    encoding = tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=256
+    )
+
+    with torch.no_grad():
+        outputs = model.generate(
+            input_ids=encoding["input_ids"],
+            attention_mask=encoding["attention_mask"],
+            max_length=256,
+            num_beams=5,
+            early_stopping=True
+        )
+
+    sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return sql.strip()
+
+# --------------------------------------------------
+# EVALUATE SQL (REWARD FUNCTION)
+# --------------------------------------------------
+def evaluate_sql(sql):
+    try:
+        cursor.execute(sql)
+        rows = cursor.fetchall()
+
+        # executed but no useful result
+        if len(rows) == 0:
+            return -0.2, rows
+
+        # good query
+        else:
+            return 1.0, rows
+
+    except Exception as e:
+        # invalid SQL
+        return -1.0, str(e)
+
+# --------------------------------------------------
+# INTERACTIVE LOOP
+# --------------------------------------------------
+while True:
+    q = input("\nAsk question (type exit to quit): ")
+
+    if q.lower() == "exit":
+        break
+
+    sql = generate_sql(q)
+
+    print("\nPredicted SQL:")
+    print(sql)
+
+    # ---------------- RUN SQL + REWARD ----------------
+    reward, output = evaluate_sql(sql)
+
+    print("\nReward:", reward)
+
+    if reward == -1.0:
+        print("SQL Error:", output)
+
+    elif reward == -0.2:
+        print("No results found")
+
+    else:
+        print("\nAnswer:")
+        for r in output:
+            print(r)

src/prepare_dataset.py

@@ -0,0 +1,143 @@
+import json
+import os
+import sqlite3
+from datasets import Dataset
+from transformers import T5Tokenizer
+
+# =========================================================
+# PROJECT ROOT (VERY IMPORTANT — fixes path issues)
+# =========================================================
+BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+TRAIN_JSON = os.path.join(BASE_DIR, "data", "train_spider.json")
+DEV_JSON   = os.path.join(BASE_DIR, "data", "dev.json")
+DB_FOLDER  = os.path.join(BASE_DIR, "data", "database")
+
+SAVE_TRAIN = os.path.join(BASE_DIR, "data", "tokenized", "train")
+SAVE_DEV   = os.path.join(BASE_DIR, "data", "tokenized", "validation")
+
+os.makedirs(os.path.dirname(SAVE_TRAIN), exist_ok=True)
+
+print("Project root:", BASE_DIR)
+print("Train file:", TRAIN_JSON)
+print("Database folder:", DB_FOLDER)
+
+# =========================================================
+# TOKENIZER
+# =========================================================
+tokenizer = T5Tokenizer.from_pretrained("t5-small")
+
+# =========================================================
+# READ DATABASE SCHEMA
+# =========================================================
+def get_schema(db_path):
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    tables = cursor.execute(
+        "SELECT name FROM sqlite_master WHERE type='table';"
+    ).fetchall()
+
+    schema_text = []
+
+    for table in tables:
+        table = table[0]
+
+        columns = cursor.execute(f"PRAGMA table_info({table});").fetchall()
+        col_names = [c[1] for c in columns]
+
+        schema_text.append(f"{table}({', '.join(col_names)})")
+
+    conn.close()
+    return "\n".join(schema_text)
+
+
+# =========================================================
+# BUILD TRAINING EXAMPLES
+# =========================================================
+def build_examples(spider_json):
+
+    print(f"\nBuilding dataset from: {spider_json}")
+
+    data = json.load(open(spider_json))
+
+    inputs = []
+    outputs = []
+
+    for ex in data:
+
+        question = ex["question"]
+        sql = ex["query"]
+        db_id = ex["db_id"]
+
+        db_path = os.path.join(DB_FOLDER, db_id, f"{db_id}.sqlite")
+
+        # skip if db missing (safety)
+        if not os.path.exists(db_path):
+            continue
+
+        schema = get_schema(db_path)
+
+        # ⭐ SCHEMA-AWARE PROMPT (VERY IMPORTANT)
+        input_text = f"""Database Schema:
+{schema}
+
+Translate English to SQL:
+{question}
+SQL:
+"""
+
+        inputs.append(input_text)
+        outputs.append(sql)
+
+    return Dataset.from_dict({"input": inputs, "output": outputs})
+
+
+# =========================================================
+# TOKENIZE
+# =========================================================
+def tokenize(example):
+
+    model_input = tokenizer(
+        example["input"],
+        max_length=512,
+        padding="max_length",
+        truncation=True
+    )
+
+    label = tokenizer(
+        example["output"],
+        max_length=256,
+        padding="max_length",
+        truncation=True
+    )
+
+    model_input["labels"] = label["input_ids"]
+    return model_input
+
+
+# =========================================================
+# RUN PIPELINE
+# =========================================================
+print("\nBuilding TRAIN dataset...")
+train_dataset = build_examples(TRAIN_JSON)
+
+print("Tokenizing TRAIN dataset...")
+tokenized_train = train_dataset.map(tokenize, batched=False)
+
+print("Saving TRAIN dataset...")
+tokenized_train.save_to_disk(SAVE_TRAIN)
+
+
+print("\nBuilding VALIDATION dataset...")
+val_dataset = build_examples(DEV_JSON)
+
+print("Tokenizing VALIDATION dataset...")
+tokenized_val = val_dataset.map(tokenize, batched=False)
+
+print("Saving VALIDATION dataset...")
+tokenized_val.save_to_disk(SAVE_DEV)
+
+print("\nDONE ✔ Dataset prepared successfully!")
+print("Train saved at:", SAVE_TRAIN)
+print("Validation saved at:", SAVE_DEV)

src/prompting.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import os
+import re
+import sqlite3
+from contextlib import closing
+from typing import Dict, Optional
+
+import torch
+
+# Keep for compatibility with existing imports. Schema linking is disabled for
+# SFT/RL alignment in this project version (full schema, deterministic order).
+USE_SCHEMA_LINKING = False
+
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DB_ROOT = os.path.join(PROJECT_ROOT, "data", "database")
+
+SCHEMA_CACHE: Dict[str, str] = {}
+
+
+def get_schema_text(db_id: str) -> str:
+    """
+    Deterministic schema string:
+      table(col1, col2, ...)
+    Tables ordered alphabetically. Columns kept in PRAGMA order.
+    """
+    if db_id in SCHEMA_CACHE:
+        return SCHEMA_CACHE[db_id]
+
+    db_path = os.path.join(DB_ROOT, db_id, f"{db_id}.sqlite")
+    schema_lines = []
+    try:
+        with closing(sqlite3.connect(db_path)) as conn:
+            cur = conn.cursor()
+            tables = cur.execute(
+                "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';"
+            ).fetchall()
+            table_names = sorted([t[0] for t in tables if t and isinstance(t[0], str)])
+            for tname in table_names:
+                cols = cur.execute(f'PRAGMA table_info("{tname}")').fetchall()
+                col_names = [c[1] for c in cols if c and isinstance(c[1], str)]
+                schema_lines.append(f"{tname}({', '.join(col_names)})")
+    except Exception:
+        schema_lines = []
+
+    schema_text = "\n".join(schema_lines).strip()
+    SCHEMA_CACHE[db_id] = schema_text
+    return schema_text
+
+
+def clean_gold_sql(sql: str) -> str:
+    """
+    Lowercase SQL + strip common Spider aliases safely.
+    If alias removal is ambiguous (same table used multiple times), keep SQL as-is.
+    """
+    if not isinstance(sql, str):
+        return ""
+    s = sql.strip().rstrip(";").strip()
+    if not s:
+        return ""
+
+    # Attempt to resolve T1/T2 aliases to table names for simple cases.
+    # Build alias -> table map from FROM/JOIN clauses.
+    alias_map: Dict[str, str] = {}
+    table_counts: Dict[str, int] = {}
+
+    for m in re.finditer(r"\b(from|join)\s+([a-zA-Z_][\w$]*)\s+(?:as\s+)?(t\d+)\b", s, flags=re.I):
+        table = m.group(2)
+        alias = m.group(3)
+        table_counts[table.lower()] = table_counts.get(table.lower(), 0) + 1
+        alias_map[alias.lower()] = table
+
+    # If any table appears multiple times, alias removal can be ambiguous → skip.
+    if any(c > 1 for c in table_counts.values()):
+        return s.lower()
+
+    # Replace alias-qualified refs alias.col -> table.col
+    out = s
+    for alias, table in alias_map.items():
+        out = re.sub(rf"\b{re.escape(alias)}\.", f"{table}.", out, flags=re.I)
+
+    # Remove alias declarations: "table AS t1" or "table t1"
+    for alias, table in alias_map.items():
+        out = re.sub(rf"\b{re.escape(table)}\s+as\s+{re.escape(alias)}\b", table, out, flags=re.I)
+        out = re.sub(rf"\b{re.escape(table)}\s+{re.escape(alias)}\b", table, out, flags=re.I)
+
+    return out.lower().strip()
+
+
+def build_prompt(
+    question: str,
+    db_id: str,
+    *,
+    schema_text: str,
+    training_sql: Optional[str] = None,
+) -> str:
+    """
+    Required prompt format:
+
+    You are a SQLite expert.
+
+    Database: <db_id>
+
+    Schema:
+    <table>(col1, col2, ...)
+    ...
+
+    Question:
+    <question>
+
+    SQL:
+    <gold sql>   (training only)
+    """
+    base = (
+        "You are a SQLite expert.\n\n"
+        f"Database: {db_id}\n\n"
+        "Schema:\n"
+        f"{schema_text}\n\n"
+        "Question:\n"
+        f"{question}\n\n"
+        "SQL:"
+    )
+    if training_sql is None:
+        return base
+    return base + "\n" + training_sql
+
+
+def encode_prompt(
+    tokenizer,
+    question: str,
+    db_id: str,
+    *,
+    device: str,
+    max_input_tokens: int = 512,
+    training_sql: Optional[str] = None,
+) -> torch.Tensor:
+    """
+    Inference mode: stops at "SQL:"
+    Training mode: can include SQL target (optional; we still recommend decoder labels).
+    Truncation happens only on schema portion by character trimming (deterministic).
+    """
+    schema_text = get_schema_text(db_id)
+    prompt = build_prompt(question, db_id, schema_text=schema_text, training_sql=training_sql)
+    enc = tokenizer(
+        prompt,
+        truncation=True,
+        max_length=max_input_tokens,
+        padding=False,
+        return_tensors="pt",
+    )
+    return enc.input_ids[0].to(device)

src/run_sql.py

@@ -0,0 +1,39 @@
+import argparse
+import json
+import os
+import sqlite3
+from contextlib import closing
+
+
+def execute_sql(db_path: str, sql: str):
+    try:
+        with closing(sqlite3.connect(db_path)) as conn:
+            cursor = conn.cursor()
+            cursor.execute(sql)
+            rows = cursor.fetchall()
+        return {"ok": True, "rows": rows, "error": None}
+    except Exception as e:
+        return {"ok": False, "rows": [], "error": str(e)}
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Safely execute SQL against a Spider SQLite DB.")
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("--db_path", type=str, help="Path to SQLite database file")
+    group.add_argument("--db_id", type=str, help="Spider database id (uses data/database/<db_id>/<db_id>.sqlite)")
+    parser.add_argument("--sql", type=str, required=True, help="SQL to execute")
+    args = parser.parse_args()
+
+    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    if args.db_path:
+        db_path = args.db_path
+    else:
+        db_path = os.path.join(project_root, "data", "database", args.db_id, f"{args.db_id}.sqlite")
+
+    result = execute_sql(db_path, args.sql)
+    print(json.dumps(result, ensure_ascii=False, default=str))
+
+
+if __name__ == "__main__":
+    main()
+

src/schema_encoder.py

@@ -0,0 +1,51 @@
+import sqlite3
+
+
+class SchemaEncoder:
+
+    def __init__(self, db_root):
+        self.db_root = db_root
+
+    def get_tables_and_columns(self, db_id):
+        db_path = self.db_root / db_id / f"{db_id}.sqlite"
+        conn = sqlite3.connect(db_path)
+        cursor = conn.cursor()
+
+        tables = cursor.execute(
+            "SELECT name FROM sqlite_master WHERE type='table';"
+        ).fetchall()
+
+        schema = {}
+
+        for (table,) in tables:
+            cols = cursor.execute(f"PRAGMA table_info({table});").fetchall()
+            col_names = [c[1] for c in cols]
+            schema[table] = col_names
+
+        conn.close()
+        return schema
+
+    # -----------------------------------
+    # Strategy 1: Structured (current)
+    # -----------------------------------
+    def structured_schema(self, db_id):
+        schema = self.get_tables_and_columns(db_id)
+
+        lines = []
+        for table, cols in schema.items():
+            lines.append(f"{table}({', '.join(cols)})")
+
+        return "\n".join(lines)
+
+    # -----------------------------------
+    # Strategy 2: Natural Language
+    # -----------------------------------
+    def natural_language_schema(self, db_id):
+        schema = self.get_tables_and_columns(db_id)
+
+        lines = []
+        for table, cols in schema.items():
+            col_text = ", ".join(cols)
+            lines.append(f"The table '{table}' contains the columns: {col_text}.")
+
+        return "\n".join(lines)

src/schema_linker.py

@@ -0,0 +1,215 @@
+"""
+Simple schema linking for Spider-style Text-to-SQL.
+
+Goal:
+- Given (question, db_id), select a small set of relevant tables/columns
+  to include in the prompt (RAG-style schema retrieval).
+
+Design constraints:
+- Pure Python (no heavy external deps).
+- Robust to missing/odd schemas: never crash.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import sqlite3
+from contextlib import closing
+from dataclasses import dataclass
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+
+_ALNUM_RE = re.compile(r"[A-Za-z0-9]+")
+_CAMEL_RE = re.compile(r"([a-z])([A-Z])")
+
+
+def _normalize_identifier(text: str) -> str:
+    """
+    Normalize a schema identifier:
+    - split underscores
+    - split camelCase / PascalCase boundaries
+    - lowercase
+    """
+    text = str(text or "")
+    text = text.replace("_", " ")
+    text = _CAMEL_RE.sub(r"\1 \2", text)
+    return text.lower()
+
+
+def _tokenize(text: str) -> List[str]:
+    text = _normalize_identifier(text)
+    return _ALNUM_RE.findall(text)
+
+
+@dataclass(frozen=True)
+class TableSchema:
+    table_name: str
+    columns: Tuple[str, ...]
+
+
+class SchemaLinker:
+    """
+    Loads Spider `tables.json` and (optionally) SQLite schemas from disk.
+    Provides a lightweight table scoring function based on token overlap.
+    """
+
+    def __init__(self, tables_json_path: str, db_root: Optional[str] = None):
+        self.tables_json_path = tables_json_path
+        self.db_root = db_root
+        self._tables_by_db: Dict[str, List[TableSchema]] = {}
+        self._sqlite_schema_cache: Dict[str, Dict[str, List[str]]] = {}
+        self._load_tables_json()
+
+    def _load_tables_json(self) -> None:
+        with open(self.tables_json_path) as f:
+            entries = json.load(f)
+
+        tables_by_db: Dict[str, List[TableSchema]] = {}
+        for entry in entries:
+            db_id = entry["db_id"]
+            table_names: List[str] = entry.get("table_names_original") or entry.get("table_names") or []
+            col_names: List[Sequence] = entry.get("column_names_original") or entry.get("column_names") or []
+
+            columns_by_table_idx: Dict[int, List[str]] = {i: [] for i in range(len(table_names))}
+            for col in col_names:
+                # Spider format: [table_idx, col_name]
+                if not col or len(col) < 2:
+                    continue
+                table_idx, col_name = col[0], col[1]
+                if table_idx is None or table_idx < 0:
+                    continue  # skip "*"
+                if table_idx not in columns_by_table_idx:
+                    continue
+                columns_by_table_idx[table_idx].append(str(col_name))
+
+            tables: List[TableSchema] = []
+            for i, tname in enumerate(table_names):
+                cols = tuple(columns_by_table_idx.get(i, []))
+                tables.append(TableSchema(table_name=str(tname), columns=cols))
+
+            tables_by_db[db_id] = tables
+
+        self._tables_by_db = tables_by_db
+
+    def _db_path(self, db_id: str) -> Optional[str]:
+        if not self.db_root:
+            return None
+        path = os.path.join(self.db_root, db_id, f"{db_id}.sqlite")
+        return path if os.path.exists(path) else None
+
+    def _load_sqlite_schema(self, db_id: str) -> Dict[str, List[str]]:
+        """
+        Load actual SQLite schema (table -> columns). Cached per db_id.
+        """
+        if db_id in self._sqlite_schema_cache:
+            return self._sqlite_schema_cache[db_id]
+
+        schema: Dict[str, List[str]] = {}
+        db_path = self._db_path(db_id)
+        if not db_path:
+            self._sqlite_schema_cache[db_id] = schema
+            return schema
+
+        try:
+            with closing(sqlite3.connect(db_path)) as conn:
+                cursor = conn.cursor()
+                tables = cursor.execute("SELECT name FROM sqlite_master WHERE type='table';").fetchall()
+                for (table_name,) in tables:
+                    columns = cursor.execute(f"PRAGMA table_info({table_name});").fetchall()
+                    schema[str(table_name)] = [str(col[1]) for col in columns]
+        except Exception:
+            schema = {}
+
+        self._sqlite_schema_cache[db_id] = schema
+        return schema
+
+    def get_schema(self, db_id: str) -> List[TableSchema]:
+        """
+        Returns a list of table schemas for this db.
+        Prefers `tables.json` (Spider canonical), but can fallback to SQLite if needed.
+        """
+        tables = self._tables_by_db.get(db_id, [])
+        if tables:
+            return tables
+
+        sqlite_schema = self._load_sqlite_schema(db_id)
+        return [TableSchema(table_name=t, columns=tuple(cols)) for t, cols in sqlite_schema.items()]
+
+    def score_tables(self, question: str, db_id: str) -> List[Tuple[float, TableSchema]]:
+        """
+        Score each table using token overlap:
+        - table token overlap (higher weight)
+        - column token overlap (lower weight)
+        """
+        q_tokens = set(_tokenize(question))
+        tables = self.get_schema(db_id)
+
+        scored: List[Tuple[float, TableSchema]] = []
+        for t in tables:
+            table_tokens = set(_tokenize(t.table_name))
+            col_tokens: set[str] = set()
+            for c in t.columns:
+                col_tokens.update(_tokenize(c))
+
+            table_overlap = len(q_tokens & table_tokens)
+            col_overlap = len(q_tokens & col_tokens)
+
+            # Simple weighted overlap (tuned to bias table matches).
+            score = 3.0 * table_overlap + 1.0 * col_overlap
+
+            # Small boost for substring mentions (helps e.g. "album" vs "albums").
+            q_text = _normalize_identifier(question)
+            if t.table_name and _normalize_identifier(t.table_name) in q_text:
+                score += 0.5
+
+            scored.append((score, t))
+
+        scored.sort(key=lambda x: (x[0], x[1].table_name), reverse=True)
+        return scored
+
+    def select_top_tables(self, question: str, db_id: str, top_k: int = 4) -> List[TableSchema]:
+        scored = self.score_tables(question, db_id)
+        if not scored:
+            return []
+        top_k = max(1, int(top_k))
+        selected = [t for _, t in scored[:top_k]]
+
+        # If everything scores 0, still return a stable selection.
+        if scored[0][0] <= 0:
+            tables = self.get_schema(db_id)
+            return tables[:top_k]
+
+        return selected
+
+    def columns_for_selected_tables(self, db_id: str, selected_tables: Iterable[TableSchema]) -> Dict[str, List[str]]:
+        """
+        Returns only columns belonging to selected tables.
+        Prefer SQLite columns (actual DB) if available; fallback to tables.json.
+        """
+        sqlite_schema = self._load_sqlite_schema(db_id)
+        out: Dict[str, List[str]] = {}
+        for t in selected_tables:
+            if t.table_name in sqlite_schema and sqlite_schema[t.table_name]:
+                out[t.table_name] = sqlite_schema[t.table_name]
+            else:
+                out[t.table_name] = list(t.columns)
+        return out
+
+    def format_relevant_schema(self, question: str, db_id: str, top_k: int = 4) -> Tuple[List[str], Dict[str, List[str]]]:
+        """
+        Returns:
+        - lines: ["table(col1, col2)", ...]
+        - selected: {table: [cols...], ...}
+        """
+        selected_tables = self.select_top_tables(question, db_id, top_k=top_k)
+        selected = self.columns_for_selected_tables(db_id, selected_tables)
+
+        lines: List[str] = []
+        for table_name, cols in selected.items():
+            cols_str = ", ".join(cols)
+            lines.append(f"{table_name}({cols_str})")
+
+        return lines, selected
+

src/sql_validator.py

@@ -0,0 +1,133 @@
+import sqlite3
+import re
+from pathlib import Path
+
+class SQLValidator:
+
+    def __init__(self, db_root):
+        self.db_root = Path(db_root)
+
+    # ---------------------------
+    # Load schema
+    # ---------------------------
+    def load_schema(self, db_id):
+        db_path = self.db_root / db_id / f"{db_id}.sqlite"
+
+        conn = sqlite3.connect(db_path)
+        cursor = conn.cursor()
+
+        tables = cursor.execute(
+            "SELECT name FROM sqlite_master WHERE type='table';"
+        ).fetchall()
+
+        schema = {}
+
+        for (table,) in tables:
+            cols = cursor.execute(f"PRAGMA table_info({table});").fetchall()
+            schema[table.lower()] = [c[1].lower() for c in cols]
+
+        conn.close()
+        return schema
+
+
+    # ---------------------------
+    # Basic syntax check
+    # ---------------------------
+    def basic_structure_valid(self, sql):
+        s = sql.lower()
+
+        if "select" not in s or "from" not in s:
+            return False, "Missing SELECT or FROM"
+
+        if len(s.split()) < 4:
+            return False, "Too short to be SQL"
+
+        return True, None
+
+
+    # ---------------------------
+    # Extract identifiers
+    # ---------------------------
+    def extract_identifiers(self, sql):
+        tokens = re.findall(r"[A-Za-z_]+", sql.lower())
+        return set(tokens)
+
+
+    # ---------------------------
+    # Table validation
+    # ---------------------------
+    def validate_tables(self, sql, schema):
+        words = self.extract_identifiers(sql)
+        tables = set(schema.keys())
+
+        used_tables = [w for w in words if w in tables]
+
+        if not used_tables:
+            return False, "No valid table used"
+
+        return True, None
+
+
+    # ---------------------------
+    # Column validation
+    # ---------------------------
+    def validate_columns(self, sql, schema):
+        words = self.extract_identifiers(sql)
+
+        valid_columns = set()
+        for cols in schema.values():
+            valid_columns.update(cols)
+
+        # ignore SQL keywords
+        keywords = {
+            "select","from","where","join","on","group","by",
+            "order","limit","count","sum","avg","min","max",
+            "and","or","in","like","distinct","asc","desc"
+        }
+
+        invalid = []
+        for w in words:
+            if w not in valid_columns and w not in schema and w not in keywords:
+                if not w.isdigit():
+                    invalid.append(w)
+
+        # allow small hallucinations but block many
+        if len(invalid) > 3:
+            return False, f"Too many unknown identifiers: {invalid[:5]}"
+
+        return True, None
+
+
+    # ---------------------------
+    # Dangerous query protection
+    # ---------------------------
+    def block_dangerous(self, sql):
+        bad = ["drop", "delete", "update", "insert", "alter"]
+
+        s = sql.lower()
+        for b in bad:
+            if b in s:
+                return False, f"Dangerous keyword detected: {b}"
+
+        return True, None
+
+
+    # ---------------------------
+    # Main validation
+    # ---------------------------
+    def validate(self, sql, db_id):
+
+        schema = self.load_schema(db_id)
+
+        checks = [
+            self.block_dangerous(sql),
+            self.basic_structure_valid(sql),
+            self.validate_tables(sql, schema),
+            self.validate_columns(sql, schema),
+        ]
+
+        for ok, msg in checks:
+            if not ok:
+                return False, msg
+
+        return True, None

src/text2sql_engine.py

@@ -0,0 +1,286 @@
+
+
+import sqlite3
+import torch
+import re
+import time
+from pathlib import Path
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+from src.sql_validator import SQLValidator
+from src.schema_encoder import SchemaEncoder  
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+DB_ROOT = PROJECT_ROOT / "data" / "database"
+
+# ==========================================
+#  UNIVERSAL STRING NORMALIZERS
+# ==========================================
+def normalize_question(q: str):
+    q = q.lower().strip()
+    q = re.sub(r"distinct\s+(\d+)", r"\1 distinct", q)
+    q = re.sub(r"\s+", " ", q)
+    return q
+
+def semantic_fix(question, sql):
+    """Universal structural fixes that apply to ALL queries and ALL databases."""
+    q = question.lower().strip()
+    s = sql.lower()
+
+    # UNIVERSAL LIMIT CATCHER: Enforce LIMIT if a number is in the question
+    #  FIXED: Removed the '?'. Now it ONLY catches numbers explicitly preceded by "show", "top", "limit", etc.
+    # This stops it from accidentally catching years like "2000".
+    num_match = re.search(r'\b(?:show|list|top|limit|get|first|last)\s+(\d+)\b', q)
+    if num_match and "limit" not in s and "count(" not in s:
+        limit_val = num_match.group(1)
+        sql = sql.rstrip(";")
+        sql = f"{sql.strip()} LIMIT {limit_val}"
+
+    return sql
+
+
+class Text2SQLEngine:
+    def __init__(self,
+                 adapter_path="checkpoints/best_rlhf_model",
+                 base_model_name="Salesforce/codet5-base",
+                 use_lora=True):
+
+        self.device = "mps" if torch.backends.mps.is_available() else (
+            "cuda" if torch.cuda.is_available() else "cpu"
+        )
+        
+        self.validator = SQLValidator(DB_ROOT)
+        self.schema_encoder = SchemaEncoder(DB_ROOT)
+        self.schema_mode = "structured"
+        
+        # Security Keywords
+        self.dml_keywords = r'\b(delete|update|insert|drop|alter|truncate)\b'
+
+        print("Loading base model...")
+        base = AutoModelForSeq2SeqLM.from_pretrained(base_model_name)
+
+        if not use_lora:
+            self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+            self.model = base.to(self.device)
+            self.model.eval()
+            print("✅ Base model ready\n")
+            return
+
+        adapter_path = (PROJECT_ROOT / adapter_path).resolve()
+        
+        print("Loading tokenizer and LoRA adapter...")
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(str(adapter_path), local_files_only=True)
+        except Exception:
+            self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+
+        self.model = PeftModel.from_pretrained(base, str(adapter_path)).to(self.device)
+        self.model.eval()
+        print("✅ RLHF model ready\n")
+
+    # ==========================================
+    # ---------------- PROMPT BUILDERS ---------
+    # ==========================================
+    def build_prompt(self, question, schema):
+        return f"""You are an expert SQL generator.
+Database schema:
+{schema}
+Generate a valid SQLite query for the question.
+Question:
+{question}
+SQL:
+"""
+
+    def build_repair_prompt(self, question, schema, bad_sql, error_msg):
+        #  UNIVERSAL UPGRADE: Extract the hallucinated column and explicitly warn the model
+        hallucinated_warning = ""
+        col_match = re.search(r"no such column:\s*([^\s]+)", error_msg, re.IGNORECASE)
+        if col_match:
+            bad_col = col_match.group(1)
+            hallucinated_warning = f"\n🚨 CRITICAL ERROR: You hallucinated the column '{bad_col}'. IT DOES NOT EXIST. Look at the schema and find the actual column name (it might be spelled differently or be a synonym like 'details', 'desc', or have a typo)."
+
+        return f"""You are an expert SQL generator.
+Database schema:
+{schema}
+
+You generated this incorrect SQL for the question "{question}":
+{bad_sql}
+
+Execution failed with this SQLite error:
+{error_msg}{hallucinated_warning}
+
+UNIVERSAL RULES TO FIX THIS:
+1. NEVER invent or guess column names. Use ONLY the exact table and column names listed in the schema above.
+2. Watch out for typos in the database schema! If you need 'assessment', look for 'asessment'. If you need 'name', look for 'details'.
+3. If the error is "no such column", you either hallucinated the name, or you forgot an INNER JOIN. Check the schema and fix it.
+4. If the query requires a COUNT() but also selects names, ensure you added a GROUP BY.
+
+Write the corrected SQLite SQL query.
+SQL:
+"""
+
+    def get_schema(self, db_id):
+        return self.schema_encoder.structured_schema(db_id)
+
+    # ==========================================
+    # ---------------- SQL POSTPROCESS ---------
+    # ==========================================
+    def extract_sql(self, text: str):
+        text = text.strip()
+        if "SQL:" in text:
+            text = text.split("SQL:")[-1]
+        match = re.search(r"select[\s\S]*", text, re.IGNORECASE)
+        if match:
+            text = match.group(0)
+        return text.split(";")[0].strip()
+
+    def clean_sql(self, sql: str):
+        sql = sql.replace('"', "'")
+        sql = re.sub(r"\s+", " ", sql)
+        return sql.strip()
+
+    def repair_logic(self, question, sql):
+        """Universal logical repairs (like missing NOT NULL for negation)"""
+        q = question.lower()
+        s = sql.lower()
+
+        # Universal Negation Auto-Joiner
+        if any(word in q for word in ["never", "no ", "without"]):
+            m = re.search(r"from\s+(\w+).*join\s+(\w+)", s)
+            if m:
+                left, right = m.group(1), m.group(2)
+                key = re.search(r"on\s+(\w+\.\w+)\s*=\s*(\w+\.\w+)", s)
+                if key:
+                    sql = f"SELECT {left}.* FROM {left} LEFT JOIN {right} ON {key.group(1)} = {key.group(2)} WHERE {key.group(2)} IS NULL"
+
+        # Universal LIKE wildcard injection
+        if any(w in q for w in ["contain", "with", "include"]):
+            sql = re.sub(r"=\s*'([^']+)'", r"LIKE '%\1%'", sql, flags=re.IGNORECASE)
+
+        return sql
+
+    # ==========================================
+    # ---------------- GENERATE ----------------
+    # ==========================================
+    def generate_sql(self, prompt, is_repair=False):
+        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(self.device)
+        
+        #  FIXED: Dynamic Generation Parameters (No more terminal warnings)
+        gen_kwargs = {
+            "max_new_tokens": 128,
+        }
+        
+        if is_repair:
+            # If the model failed, it needs to think differently. 
+            # We turn off rigid beam search and introduce sampling so it doesn't repeat the exact same broken SQL.
+            gen_kwargs["do_sample"] = True
+            gen_kwargs["temperature"] = 0.5
+            gen_kwargs["top_p"] = 0.9
+        else:
+            # First attempt is strictly deterministic for maximum benchmark accuracy
+            gen_kwargs["num_beams"] = 5
+            gen_kwargs["do_sample"] = False
+            gen_kwargs["early_stopping"] = True # <--- Moved here so it doesn't clash with sampling!
+
+        with torch.no_grad():
+            outputs = self.model.generate(**inputs, **gen_kwargs)
+            
+        decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return self.clean_sql(self.extract_sql(decoded))
+
+    # ==========================================
+    # ---------------- EXECUTE -----------------
+    # ==========================================
+    def execute_sql(self, question, sql, db_id):
+        
+        # 🛡️ DEFENSE LAYER 2: Block Execution of Malicious SQL
+        if re.search(self.dml_keywords, sql, re.IGNORECASE):
+            return sql, [], [], "❌ Security Alert: Malicious DML/DDL SQL syntax blocked."
+
+        db_path = DB_ROOT / db_id / f"{db_id}.sqlite"
+
+        sql = self.repair_logic(question, sql)
+        sql = self.clean_sql(sql)
+        sql = semantic_fix(question, sql)
+
+        is_valid, reason = self.validator.validate(sql, db_id)
+        if not is_valid:
+            return sql, [], [], f"Blocked unsafe SQL: {reason}"
+
+        try:
+            conn = sqlite3.connect(db_path)
+            start_time = time.monotonic()
+            def timeout_handler():
+                return 1 if (time.monotonic() - start_time) > 5.0 else 0
+            conn.set_progress_handler(timeout_handler, 10000)
+
+            cursor = conn.cursor()
+            cursor.execute(sql)
+            rows = cursor.fetchall()
+            columns = [d[0] for d in cursor.description] if cursor.description else []
+            conn.close()
+            
+            return sql, columns, rows, None
+
+        except Exception as e:
+            return sql, [], [], str(e)
+
+    # ==========================================
+    # ---------------- PIPELINE ----------------
+    # ==========================================
+    def ask(self, question, db_id):
+        question = normalize_question(question)
+        
+        # 🛡️ DEFENSE LAYER 1: Block Malicious Natural Language Intent Early
+        if re.search(self.dml_keywords, question, re.IGNORECASE):
+            return {
+                "question": question,
+                "sql": "-- BLOCKED",
+                "columns": [],
+                "rows": [],
+                "error": "❌ Security Alert: Malicious intent (DELETE/DROP/UPDATE) detected in the prompt."
+            }
+
+        # 1. First Pass Generation
+        schema = self.get_schema(db_id)
+        prompt = self.build_prompt(question, schema)
+        
+        # is_repair=False -> Uses strict Beam Search
+        raw_sql = self.generate_sql(prompt, is_repair=False) 
+        
+        # 2. First Execution Attempt
+        final_sql, cols, rows, error = self.execute_sql(question, raw_sql, db_id)
+        
+        # 🤖 3. UNIVERSAL AGENTIC SELF-CORRECTION LOOP
+        if error and "Security Alert" not in error:
+            print(f"\n Caught SQLite Error: {error}")
+            print(f" Triggering Stochastic LLM Self-Correction...")
+            
+            # Feed the explicit error instructions back to the LLM
+            repair_prompt = self.build_repair_prompt(question, schema, final_sql, error)
+            
+            # is_repair=True -> Uses Temperature Sampling to break out of hallucination loops
+            repaired_sql = self.generate_sql(repair_prompt, is_repair=True) 
+            
+            # Try executing the repaired SQL
+            final_sql, cols, rows, error = self.execute_sql(question, repaired_sql, db_id)
+            
+            if not error:
+                print("✅ Universal Agent successfully self-corrected the query!")
+            else:
+                print("❌ Model failed self-correction.")
+        
+        return {
+            "question": question, 
+            "sql": final_sql, 
+            "columns": cols, 
+            "rows": rows, 
+            "error": error
+        }
+
+_engine = None
+def get_engine():
+    global _engine
+    if _engine is None:
+        _engine = Text2SQLEngine()
+    return _engine

src/tokenize_dataset.py

@@ -0,0 +1,57 @@
+from datasets import Dataset
+from transformers import T5Tokenizer
+import pandas as pd
+
+print("Loading processed dataset...")
+train = pd.read_csv("../data/processed/train.csv")
+val   = pd.read_csv("../data/processed/validation.csv")
+
+# remove hidden pandas index column if exists
+train = train.drop(columns=[c for c in train.columns if "index" in c.lower()], errors="ignore")
+val   = val.drop(columns=[c for c in val.columns if "index" in c.lower()], errors="ignore")
+
+print("Loading tokenizer (t5-small)...")
+tokenizer = T5Tokenizer.from_pretrained("t5-small")
+
+SQL_PREFIX = "translate English to SQL: "
+
+# ------------------------------------------------------
+# TOKENIZATION FUNCTION
+# ------------------------------------------------------
+def tokenize(example):
+
+    # input = schema + question
+    input_text = SQL_PREFIX + example["input"]
+
+    # target = real SQL
+    target_sql = example["sql"]
+
+    model_inputs = tokenizer(
+        input_text,
+        text_target=target_sql,
+        max_length=256,
+        padding="max_length",
+        truncation=True
+    )
+
+    return model_inputs
+
+
+# ------------------------------------------------------
+# DATASET CONVERSION
+# ------------------------------------------------------
+print("Preparing dataset...")
+train_ds = Dataset.from_pandas(train)
+val_ds   = Dataset.from_pandas(val)
+
+print("Tokenizing train...")
+train_ds = train_ds.map(tokenize, remove_columns=train_ds.column_names)
+
+print("Tokenizing validation...")
+val_ds = val_ds.map(tokenize, remove_columns=val_ds.column_names)
+
+# save tokenized dataset
+train_ds.save_to_disk("../data/tokenized/train")
+val_ds.save_to_disk("../data/tokenized/validation")
+
+print("DONE ✔ Tokenized dataset saved correctly")

src/train_model.py

@@ -0,0 +1,105 @@
+import torch
+from datasets import load_from_disk
+from transformers import (
+    T5ForConditionalGeneration,
+    T5Tokenizer,
+    DataCollatorForSeq2Seq,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments
+)
+
+# ======================================================
+# DEVICE (Mac M1/M2/M3 Safe)
+# ======================================================
+device = "mps" if torch.backends.mps.is_available() else "cpu"
+print("Using device:", device)
+
+# ======================================================
+# LOAD TOKENIZED DATASET (FIXED PATHS)
+# ======================================================
+print("Loading tokenized dataset...")
+
+train_dataset = load_from_disk("data/tokenized/train")
+val_dataset   = load_from_disk("data/tokenized/validation")
+
+print("Train size:", len(train_dataset))
+print("Validation size:", len(val_dataset))
+
+# ======================================================
+# LOAD MODEL
+# ======================================================
+print("Loading model (t5-small)...")
+
+model = T5ForConditionalGeneration.from_pretrained("t5-small").to(device)
+tokenizer = T5Tokenizer.from_pretrained("t5-small")
+
+# Prevent Mac memory crash
+model.config.use_cache = False
+
+# Important T5 settings (prevents generation bugs)
+model.config.decoder_start_token_id = tokenizer.pad_token_id
+model.config.eos_token_id = tokenizer.eos_token_id
+model.config.pad_token_id = tokenizer.pad_token_id
+
+# ======================================================
+# DATA COLLATOR
+# ======================================================
+data_collator = DataCollatorForSeq2Seq(
+    tokenizer=tokenizer,
+    model=model
+)
+
+# ======================================================
+# TRAINING ARGUMENTS (Mac Safe)
+# ======================================================
+print("Setting training config...")
+
+training_args = Seq2SeqTrainingArguments(
+    output_dir="outputs/model",
+
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+
+    learning_rate=3e-4,
+    num_train_epochs=5,
+
+    per_device_train_batch_size=1,
+    per_device_eval_batch_size=1,
+    gradient_accumulation_steps=8,
+
+    logging_steps=50,
+
+    fp16=False,
+    bf16=False,
+    dataloader_pin_memory=False,
+
+    predict_with_generate=True,
+    report_to="none"
+)
+
+# ======================================================
+# TRAINER
+# ======================================================
+trainer = Seq2SeqTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=val_dataset,
+    tokenizer=tokenizer,
+    data_collator=data_collator,
+)
+
+# ======================================================
+# TRAIN
+# ======================================================
+print("Training started 🚀")
+trainer.train()
+
+# ======================================================
+# SAVE MODEL
+# ======================================================
+print("Saving model...")
+trainer.save_model("outputs/model")
+tokenizer.save_pretrained("outputs/model")
+
+print("\nDONE ✔ Base model trained successfully")

src/train_rl.py

@@ -0,0 +1,816 @@
+# =========================================================
+# RLHF TRAINING FOR TEXT2SQL (STABLE PPO VERSION)
+# =========================================================
+
+import torch
+import torch.nn.functional as F
+from datasets import load_dataset
+from transformers import AutoTokenizer
+from transformers.generation.logits_process import LogitsProcessor, LogitsProcessorList
+from trl import PPOTrainer, PPOConfig, AutoModelForSeq2SeqLMWithValueHead
+from peft import PeftModel
+import os, sys, sqlite3, re, random
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+from execution_reward import execution_reward, extract_tables, extract_columns
+try:
+    import sqlparse  # gate PPO updates on parsable SQL only
+except Exception:  # pragma: no cover
+    sqlparse = None
+
+
+# ======================================================
+# DEVICE
+# ======================================================
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
+print("Using device:", device)
+
+
+# ======================================================
+# TRAINING SETTINGS
+# ======================================================
+NUM_EPOCHS = 5
+LOG_EVERY = 20
+USE_SCHEMA = True
+SCHEMA_WARMUP_EPOCHS = 0
+MAX_SCHEMA_CHARS = 1500
+MAX_OUTPUT_TOKENS = 80
+ROLLOUTS_PER_EPOCH = 2048
+
+
+# ======================================================
+# PATHS
+# ======================================================
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+# 🎯 FIXED: Save ONLY the best model to this exact path
+RL_MODEL_PATH = os.path.join(PROJECT_ROOT, "checkpoints", "rlhf_t5_best")
+output_dir = RL_MODEL_PATH
+
+DB_ROOT = os.path.join(PROJECT_ROOT, "data/database")
+
+# 🎯 Updated to point to our newly trained t5-small SFT model
+ADAPTER_PATH = os.path.join(PROJECT_ROOT, "checkpoints/sft_t5") 
+
+FALLBACK_ADAPTER_PATH = os.path.join(PROJECT_ROOT, "models/t5_spider_sft_lora")
+FALLBACK_ADAPTER_PATH_2 = os.path.join(PROJECT_ROOT, "outputs/sft_text2sql")
+# 🎯 ENSURING t5-small is used
+BASE_MODEL = os.environ.get("BASE_MODEL", "t5-small")
+
+
+# ======================================================
+# LOAD MODEL (LoRA)
+# ======================================================
+print("Loading base:", BASE_MODEL)
+if not os.path.isdir(ADAPTER_PATH):
+    if os.path.isdir(FALLBACK_ADAPTER_PATH):
+        ADAPTER_PATH = FALLBACK_ADAPTER_PATH
+    elif os.path.isdir(FALLBACK_ADAPTER_PATH_2):
+        ADAPTER_PATH = FALLBACK_ADAPTER_PATH_2
+print("Loading adapters:", ADAPTER_PATH)
+
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+
+model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(BASE_MODEL).to(device)
+model.pretrained_model = PeftModel.from_pretrained(model.pretrained_model, ADAPTER_PATH)
+
+ref_model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(BASE_MODEL).to(device)
+ref_model.pretrained_model = PeftModel.from_pretrained(ref_model.pretrained_model, ADAPTER_PATH)
+
+ref_model.eval()
+for p in ref_model.parameters():
+    p.requires_grad_(False)
+
+# Freeze base transformer weights; train LoRA adapters + value head.
+for name, p in model.named_parameters():
+    # Train value head
+    if name.startswith("v_head"):
+        p.requires_grad = True
+    # Train LoRA adapters (policy learning!)
+    elif "lora_" in name:
+        p.requires_grad = True
+    # Freeze base model
+    else:
+        p.requires_grad = False
+
+trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+total = sum(p.numel() for p in model.parameters())
+print(f"Trainable params: {trainable}/{total} ({100*trainable/total:.2f}%)")
+
+model.config.use_cache = False
+ref_model.config.use_cache = False
+
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+
+# ======================================================
+# DATASET
+# ======================================================
+print("Loading Spider subset...")
+random.seed(0)
+
+# Train on a small, stable curriculum of DBs first.
+TRAIN_DBS = [
+    "flight_1",
+    "student_assessment",
+    "store_1",
+    "bike_1",
+    "book_2",
+    "chinook_1",
+]
+
+dataset = load_dataset("spider", split="train")
+_TRAIN_DBS_SET = set(TRAIN_DBS)
+dataset = dataset.filter(lambda x: x["db_id"] in _TRAIN_DBS_SET)
+dataset = dataset.select(range(min(800, len(dataset))))
+
+print("Using RLHF DBs:", TRAIN_DBS)
+print("Filtered size:", len(dataset))
+
+total_steps = ROLLOUTS_PER_EPOCH
+
+# ======================================================
+# DB UTILITIES
+# ======================================================
+def get_db_path(db_id):
+    return os.path.join(DB_ROOT, db_id, f"{db_id}.sqlite")
+
+
+def get_db_schema(db_path):
+    schema_text = ""
+    try:
+        conn = sqlite3.connect(db_path)
+        cursor = conn.cursor()
+
+        tables = cursor.execute(
+            "SELECT name FROM sqlite_master WHERE type='table';"
+        ).fetchall()
+
+        for table in tables:
+            table_name = table[0]
+            columns = cursor.execute(f"PRAGMA table_info({table_name});").fetchall()
+            col_names = [col[1] for col in columns]
+            schema_text += f"{table_name}({', '.join(col_names)}) "
+
+        conn.close()
+    except:
+        pass
+
+    return schema_text
+
+
+# ======================================================
+# PROMPT
+# ======================================================
+PREFIX = "translate English to SQL:"
+
+
+def trim_schema(schema: str, max_chars: int = 1200) -> str:
+    if schema is None:
+        return ""
+    schema = str(schema)
+    if len(schema) <= max_chars:
+        return schema
+    return schema[:max_chars]
+def build_prompt(question: str, schema: str, use_schema: bool) -> str:
+    if not use_schema:
+        return f"{PREFIX}\n\nQuestion:\n{question}\n\nSQL:"
+
+    schema = trim_schema(schema, max_chars=MAX_SCHEMA_CHARS)
+    return f"{PREFIX}\n\nSchema:\n{schema}\n\nQuestion:\n{question}\n\nSQL:"
+
+def encode_prompt(question, schema, use_schema):
+    # Never truncate the question; only truncate schema tokens if needed.
+    if not use_schema:
+        prompt = build_prompt(question, schema, use_schema=False)
+        return tokenizer(prompt, return_tensors="pt", truncation=True).input_ids[0].to(device)
+
+    schema = trim_schema(schema, max_chars=MAX_SCHEMA_CHARS)
+    prefix_schema = f"{PREFIX}\n\nSchema:\n"
+    mid = "\n\nQuestion:\n"
+    suffix = f"{question}\n\nSQL:"
+
+    prefix_ids = tokenizer.encode(prefix_schema, add_special_tokens=False)
+    schema_ids = tokenizer.encode(schema, add_special_tokens=False)
+    mid_ids = tokenizer.encode(mid, add_special_tokens=False)
+    suffix_ids = tokenizer.encode(suffix, add_special_tokens=False)
+
+    max_len = getattr(tokenizer, "model_max_length", 512)
+    eos_id = tokenizer.eos_token_id
+    max_without_eos = max_len - (1 if eos_id is not None else 0)
+
+    # Ensure the question+SQL suffix always fits; truncate schema first.
+    fixed_len = len(prefix_ids) + len(mid_ids) + len(suffix_ids)
+    if fixed_len > max_without_eos:
+        # Extremely rare; clip the suffix (question) only if unavoidable.
+        keep = max(0, max_without_eos - (len(prefix_ids) + len(mid_ids)))
+        suffix_ids = suffix_ids[:keep]
+        fixed_len = len(prefix_ids) + len(mid_ids) + len(suffix_ids)
+
+    remaining_for_schema = max_without_eos - fixed_len
+    if remaining_for_schema < 0:
+        remaining_for_schema = 0
+    schema_ids = schema_ids[:remaining_for_schema]
+
+    ids = prefix_ids + schema_ids + mid_ids + suffix_ids
+    ids = ids[:max_without_eos]
+    if eos_id is not None:
+        ids = ids + [eos_id]
+
+    return torch.tensor(ids, dtype=torch.long).to(device)
+
+
+# ======================================================
+# SQL CONSTRAINED DECODING
+# ======================================================
+SQL_KEYWORDS = [
+    "select", "from", "where", "join", "inner", "left", "right",
+    "full", "outer", "on", "group", "by", "order", "having",
+    "limit", "distinct", "as", "and", "or", "not", "in", "is",
+    "null", "like", "between", "asc", "desc", "union",
+    "intersect", "except",
+]
+
+SQL_OPERATORS = ["*", ",", ".", "(", ")", "=", "<", ">", "!", "+", "-", "/", "%", "_"]
+
+
+def _piece_token_str(tok: str) -> str:
+    # T5 SentencePiece: "▁" marks a leading space; strip it for char checks.
+    return tok.lstrip("▁")
+
+
+def _precompute_always_allowed_token_ids():
+    vocab_size = len(tokenizer)
+    allowed = set()
+
+    # Always allow special tokens.
+    for tid in [tokenizer.pad_token_id, tokenizer.eos_token_id, tokenizer.unk_token_id]:
+        if tid is not None and tid >= 0:
+            allowed.add(int(tid))
+
+    # Allow whitespace/newlines in case they exist as pieces.
+    for s in [" ", "\n", "\t"]:
+        allowed.update(tokenizer.encode(s, add_special_tokens=False))
+
+    # Allow operators/punctuation/numeric pieces broadly.
+    op_chars = set("".join(SQL_OPERATORS))
+    for tid in range(vocab_size):
+        tok = tokenizer.convert_ids_to_tokens(tid)
+        if not isinstance(tok, str) or not tok:
+            continue
+        piece = _piece_token_str(tok)
+        if not piece:
+            continue
+        if all((ch in op_chars) for ch in piece):
+            allowed.add(tid)
+            continue
+        if piece.isdigit():
+            allowed.add(tid)
+            continue
+        # Common numeric fragments like "1", "00", etc.
+        if all(ch.isdigit() for ch in piece):
+            allowed.add(tid)
+
+    # Allow keyword pieces.
+    for kw in SQL_KEYWORDS:
+        for variant in {kw, kw.upper(), kw.capitalize()}:
+            allowed.update(tokenizer.encode(" " + variant, add_special_tokens=False))
+            allowed.update(tokenizer.encode(variant, add_special_tokens=False))
+
+    return allowed
+
+
+ALWAYS_ALLOWED_TOKEN_IDS = _precompute_always_allowed_token_ids()
+
+
+def _schema_allowed_token_ids(table_names, column_names):
+    allowed = set(ALWAYS_ALLOWED_TOKEN_IDS)
+
+    def _add_identifier(name: str):
+        if not name:
+            return
+        # Add whole identifier and common splits.
+        variants = {name, name.lower(), name.upper()}
+        parts = re.split(r"[_\s]+", name)
+        variants.update({p for p in parts if p})
+        for v in variants:
+            allowed.update(tokenizer.encode(" " + v, add_special_tokens=False))
+            allowed.update(tokenizer.encode(v, add_special_tokens=False))
+
+    for t in table_names:
+        _add_identifier(t)
+    for c in column_names:
+        _add_identifier(c)
+
+    return allowed
+
+
+class SQLVocabularyLogitsProcessor(LogitsProcessor):
+    def __init__(self, allowed_token_ids):
+        self.allowed_token_ids = {int(i) for i in allowed_token_ids if int(i) >= 0}
+        self._bias = None
+        self._bias_vocab_size = None
+
+    def _get_bias(self, scores: torch.Tensor) -> torch.Tensor:
+        vocab_size = int(scores.shape[-1])
+        if (
+            self._bias is None
+            or self._bias.device != scores.device
+            or self._bias.dtype != scores.dtype
+            or self._bias_vocab_size != vocab_size
+        ):
+            bias = torch.full((vocab_size,), float("-inf"), device=scores.device, dtype=scores.dtype)
+            for tid in self.allowed_token_ids:
+                if tid < vocab_size:
+                    bias[tid] = 0.0
+            self._bias = bias
+            self._bias_vocab_size = vocab_size
+        return self._bias
+
+    def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
+        return scores + self._get_bias(scores)
+
+
+_DB_VOCAB_CACHE = {}
+
+
+def get_db_tables_columns(db_path: str):
+    if db_path in _DB_VOCAB_CACHE:
+        return _DB_VOCAB_CACHE[db_path]
+    tables, cols = [], []
+    try:
+        conn = sqlite3.connect(db_path)
+        cur = conn.cursor()
+        for (tname,) in cur.execute(
+            "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';"
+        ).fetchall():
+            if not tname:
+                continue
+            tables.append(tname)
+            try:
+                for row in cur.execute(f'PRAGMA table_info("{tname}")').fetchall():
+                    if row and isinstance(row[1], str):
+                        cols.append(row[1])
+            except Exception:
+                continue
+        conn.close()
+    except Exception:
+        pass
+    _DB_VOCAB_CACHE[db_path] = (tables, cols)
+    return tables, cols
+
+
+# ======================================================
+# PPO CONFIG (stable learning)
+# ======================================================
+ppo_config = PPOConfig(
+    learning_rate=2e-5,            # was 1e-6 → model could not move
+    batch_size=8,                  # better gradient estimate
+    mini_batch_size=2,
+    gradient_accumulation_steps=2, # stable updates on small data
+    ppo_epochs=1,
+
+    # --- KL control (MOST IMPORTANT FIX) ---
+    init_kl_coef=0.05,             # reduce punishment
+    target_kl=0.15,                # relax constraint to avoid skipped updates
+    adap_kl_ctrl=True,
+
+    # --- stability ---
+    cliprange=0.25,
+    cliprange_value=0.25,
+    whiten_rewards=True,
+    kl_penalty="kl",
+    max_grad_norm=1.0,
+)
+trainer = PPOTrainer(
+    config=ppo_config,
+    model=model,
+    ref_model=ref_model,
+    tokenizer=tokenizer,
+)
+optimizer = trainer.optimizer
+
+# Provide `.device` attribute for the supervised anchor helper.
+try:
+    model.device = torch.device(device)
+except Exception:
+    pass
+
+
+# ======================================================
+# GENERATION (schema-constrained decoding)
+# ======================================================
+generation_kwargs = dict(
+    max_new_tokens=64,         # 128 causes garbage SQL loops
+
+    do_sample=True,
+    temperature=0.9,           # encourage exploration
+    top_p=0.95,
+    top_k=100,
+
+    repetition_penalty=1.1,    # prevents SELECT SELECT SELECT
+    no_repeat_ngram_size=3,
+
+    num_beams=1,
+    pad_token_id=tokenizer.pad_token_id,
+    eos_token_id=tokenizer.eos_token_id,
+)
+# ======================================================
+# TRAIN LOOP
+# ======================================================
+print("Starting RL training 🚀")
+
+query_buffer, response_buffer, reward_buffer, gold_buffer = [], [], [], []
+query_text_buffer = []
+
+best_reward = -999999
+best_epoch = -1
+
+def _is_parsable_sql(sql: str) -> bool:
+    s = (sql or "").strip()
+    if not s:
+        return False
+    up = s.upper()
+    if "SELECT" not in up or "FROM" not in up:
+        return False
+    if sqlparse is None:
+        return True
+    try:
+        return bool(sqlparse.parse(s))
+    except Exception:
+        return False
+
+
+def _pad_2d(seqs, pad_id: int):
+    max_len = max(int(s.numel()) for s in seqs)
+    out = torch.full((len(seqs), max_len), int(pad_id), dtype=torch.long, device=device)
+    attn = torch.zeros((len(seqs), max_len), dtype=torch.long, device=device)
+    for i, s in enumerate(seqs):
+        n = int(s.numel())
+        out[i, :n] = s.to(device)
+        attn[i, :n] = 1
+    return out, attn
+
+
+def _shift_right(labels: torch.Tensor, start_id: int) -> torch.Tensor:
+    dec = labels.clone()
+    dec[:, 1:] = labels[:, :-1]
+    dec[:, 0] = int(start_id)
+    return dec
+
+
+def safe_get_kl(stats):
+    if not isinstance(stats, dict):
+        return None
+    for k in stats.keys():
+        if "kl" in str(k).lower():
+            v = stats[k]
+            try:
+                return float(v.item() if hasattr(v, "item") else v)
+            except Exception:
+                return None
+    return None
+
+def supervised_anchor_step(model, tokenizer, queries, gold_sqls, weight=0.05):
+    model.train()
+    total_loss = 0.0
+
+    for q, gold in zip(queries, gold_sqls):
+
+        enc = tokenizer(q, return_tensors="pt", truncation=True).to(model.device)
+        dec = tokenizer(text_target=gold, return_tensors="pt", truncation=True)
+
+        labels = dec.input_ids.to(model.device)
+
+        # teacher forcing shift
+        decoder_input_ids = labels[:, :-1].contiguous()
+        target_ids = labels[:, 1:].contiguous()
+
+        outputs = model(
+            input_ids=enc.input_ids,
+            attention_mask=enc.attention_mask,
+            decoder_input_ids=decoder_input_ids,
+        )
+
+        logits = outputs[0]
+
+        vocab_size = logits.size(-1)
+        loss = F.cross_entropy(
+            logits.view(-1, vocab_size),
+            target_ids.view(-1),
+            ignore_index=tokenizer.pad_token_id,
+        )
+
+        (loss * weight).backward()
+        total_loss += loss.item()
+
+    return total_loss
+
+
+@torch.no_grad()
+def _estimate_policy_entropy(query_tensors, response_tensors) -> torch.Tensor:
+    """
+    Returns per-sample average token entropy of the policy on the sampled response tokens.
+    Used as a small bonus to reduce repetition collapse.
+    """
+    pad_id = int(tokenizer.pad_token_id)
+    enc_ids, enc_attn = _pad_2d(query_tensors, pad_id)
+    dec_ids, dec_attn = _pad_2d(response_tensors, pad_id)
+
+    start_id = int(getattr(model.pretrained_model.config, "decoder_start_token_id", pad_id))
+    dec_inp = _shift_right(dec_ids, start_id)
+
+    out = model.pretrained_model(
+        input_ids=enc_ids,
+        attention_mask=enc_attn,
+        decoder_input_ids=dec_inp,
+        use_cache=False,
+    )
+    logp = torch.log_softmax(out.logits, dim=-1)
+    p = torch.exp(logp)
+    ent = -(p * logp).sum(dim=-1)  # [B, T]
+    # average only over non-pad positions of the sampled response
+    denom = dec_attn.sum(dim=-1).clamp_min(1)
+    return (ent * dec_attn).sum(dim=-1) / denom  # [B]
+
+
+def _repeat_penalty(response_tensor: torch.Tensor) -> float:
+    """
+    Penalize repetition to avoid 'SELECT SELECT SELECT' collapse.
+    Simple heuristic: consecutive duplicate token ratio + low-unique-token ratio.
+    """
+    ids = response_tensor.detach().tolist()
+    n = len(ids)
+    if n <= 1:
+        return 0.0
+    consec_dup = 0
+    for i in range(1, n):
+        if ids[i] == ids[i - 1]:
+            consec_dup += 1
+    unique_ratio = len(set(ids)) / n
+    consec_ratio = consec_dup / (n - 1)
+    # Higher penalty when low unique + high consecutive duplicates
+    return float(0.5 * consec_ratio + 0.5 * (1.0 - unique_ratio))
+
+
+def _supervised_anchor_step(query_tensors, gold_sql_texts, weight: float = 0.05) -> None:
+    """
+    Small teacher-forcing step on gold SQL to anchor grammar during PPO.
+    Runs only if PPOTrainer exposes (accelerator, optimizer).
+    """
+    if not gold_sql_texts:
+        return
+    accelerator = getattr(trainer, "accelerator", None)
+    optimizer = getattr(trainer, "optimizer", None)
+    if accelerator is None or optimizer is None:
+        return
+
+    pad_id = int(tokenizer.pad_token_id)
+    enc_ids, enc_attn = _pad_2d(query_tensors, pad_id)
+
+    # Tokenize gold SQL targets (decoder side)
+    gold_ids = []
+    for s in gold_sql_texts:
+        g = (s or "").strip()
+        if not g:
+            g = "SELECT 1"
+        ids = tokenizer.encode(g, add_special_tokens=False)[:256]
+        if tokenizer.eos_token_id is not None:
+            ids = ids + [int(tokenizer.eos_token_id)]
+        gold_ids.append(torch.tensor(ids, dtype=torch.long))
+
+    dec_ids, dec_attn = _pad_2d(gold_ids, pad_id)
+    labels = dec_ids.clone()
+    labels[dec_attn == 0] = -100
+
+    # PEFT model forward supports labels -> returns loss
+    out = model.pretrained_model(
+        input_ids=enc_ids,
+        attention_mask=enc_attn,
+        labels=labels,
+        use_cache=False,
+    )
+    loss = out.loss * float(weight)
+
+    optimizer.zero_grad(set_to_none=True) if hasattr(optimizer, "zero_grad") else None
+    accelerator.backward(loss)
+    optimizer.step()
+
+
+def _curriculum_allows(gold_sql: str, epoch_num: int) -> bool:
+    gold_up = (gold_sql or "").upper()
+    has_join = "JOIN" in gold_up
+    has_set_op = any(op in gold_up for op in ["UNION", "INTERSECT", "EXCEPT"])
+    tables = extract_tables(gold_sql)
+    single_table = len(tables) <= 1 and (not has_join)
+
+    # Epoch 1: only single-table, no joins/set-ops.
+    if epoch_num == 1:
+        return single_table and (not has_set_op)
+    # Epoch 2: allow joins, but still avoid set-ops.
+    if epoch_num == 2:
+        return (single_table or has_join) and (not has_set_op)
+    # Epoch 3+: full dataset.
+    return True
+
+
+for epoch in range(1, NUM_EPOCHS + 1):
+
+    use_schema_this_epoch = USE_SCHEMA and (epoch > SCHEMA_WARMUP_EPOCHS)
+
+    epoch_reward_sum = 0
+    negative_rewards = 0
+    partial_rewards = 0
+    correct_rewards = 0
+
+    total_considered = 0
+    valid_sql_count = 0
+    exec_correct_count = 0
+    table_overlap_sum = 0.0
+    column_overlap_sum = 0.0
+    kl_values = []
+
+    for step in range(1, total_steps + 1):
+
+        example = dataset[random.randrange(len(dataset))]
+
+        question = example["question"]
+        gold_sql = example["query"]
+        db_id = example["db_id"]
+        db_path = get_db_path(db_id)
+
+        # NOTE: sampling-with-replacement provides more rollouts per epoch.
+
+        schema = get_db_schema(db_path)
+        question_text = build_prompt(question, schema, use_schema_this_epoch)
+        query_tensor = encode_prompt(question, schema, use_schema_this_epoch)
+
+        # ----- generate -----
+        table_names, column_names = get_db_tables_columns(db_path)
+        allowed_ids = _schema_allowed_token_ids(table_names, column_names)
+        logits_processor = LogitsProcessorList([SQLVocabularyLogitsProcessor(allowed_ids)])
+
+        response = trainer.generate([query_tensor], logits_processor=logits_processor, **generation_kwargs)[0]
+        response_tensor = response.squeeze(0)[:MAX_OUTPUT_TOKENS]
+
+        pred_sql = tokenizer.decode(response_tensor.cpu(), skip_special_tokens=True)
+
+        total_considered += 1
+
+        # PPO must optimize ONLY when SQL parses successfully.
+        if not _is_parsable_sql(pred_sql):
+            negative_rewards += 1
+            continue
+
+        # Reject generations shorter than 6 tokens.
+        if int(response_tensor.numel()) < 6:
+            negative_rewards += 1
+            continue
+
+        # ----- reward -----
+        reward_value = execution_reward(pred_sql, db_path, gold_sql)
+
+        # SQL validity gate: if invalid/unparsable -> reward_value is None -> skip PPO entirely.
+        if reward_value is None:
+            if step % 100 == 0:
+                ratio = valid_sql_count / max(total_considered, 1)
+                print(f"\nLearning ratio: {valid_sql_count}/{total_considered} ({ratio:.3f})")
+                if ratio < 0.15:
+                    print("MODEL COLLAPSING")
+            continue
+
+        # Clip rewards to [-1, 1]
+        reward_value = float(max(-1.0, min(1.0, reward_value)))
+        # Penalize repetition in decoded output (token-level heuristic).
+        reward_value = float(max(-1.0, min(1.0, reward_value - 0.2 * _repeat_penalty(response_tensor))))
+        # Keep rewards on CPU for normalization; move to device only for trainer.step().
+        reward_tensor = torch.tensor(reward_value, dtype=torch.float32)
+
+        epoch_reward_sum += reward_value
+
+        # ----- metrics -----
+        # "Valid sample" means reward is not None (parsable SQL).
+        valid_sql_count += 1
+
+        pred_tables = extract_tables(pred_sql)
+        gold_tables = extract_tables(gold_sql)
+        pred_cols = extract_columns(pred_sql)
+        gold_cols = extract_columns(gold_sql)
+
+        if len(gold_tables) > 0:
+            table_overlap_sum += len(pred_tables & gold_tables) / max(len(gold_tables), 1)
+        if len(gold_cols) > 0:
+            column_overlap_sum += len(pred_cols & gold_cols) / max(len(gold_cols), 1)
+
+        # execution_reward returns 1.0 for correct execution result.
+        if reward_value >= 1.0:
+            exec_correct_count += 1
+
+        if reward_value <= -1.0:
+            negative_rewards += 1
+        elif reward_value >= 1.0:
+            correct_rewards += 1
+        else:
+            partial_rewards += 1
+
+        # Train only on informative samples:
+        # - invalid SQL already skipped (reward is None)
+        # - very small magnitude signal skipped
+        if abs(reward_value) < 0.1:
+            continue
+
+        query_buffer.append(query_tensor)
+        response_buffer.append(response_tensor)
+        reward_buffer.append(reward_tensor)
+        gold_buffer.append(gold_sql)
+        query_text_buffer.append(question_text)
+
+        # ----- PPO update -----
+        if len(query_buffer) == ppo_config.batch_size:
+            # move rewards to device
+            reward_buffer = [r.to(device) for r in reward_buffer]
+
+            # run PPO step
+            stats = trainer.step(query_buffer, response_buffer, reward_buffer)
+
+            # log KL safely (no control logic)
+            kl = safe_get_kl(stats)
+            if kl is not None:
+                kl_values.append(kl)
+
+            # --- supervised anchor to prevent grammar collapse ---
+            supervised_anchor_step(model, tokenizer, query_text_buffer, gold_buffer, weight=0.05)
+            optimizer.step()
+            optimizer.zero_grad()
+
+            # reset buffers
+            query_buffer, response_buffer, reward_buffer, gold_buffer = [], [], [], []
+            query_text_buffer = []
+
+        # ----- learning ratio logging -----
+        if step % 100 == 0:
+            ratio = valid_sql_count / max(total_considered, 1)
+            print(f"\nLearning ratio: {valid_sql_count}/{total_considered} ({ratio:.3f})")
+            if ratio < 0.15:
+                print("MODEL COLLAPSING")
+                # Increase KL coefficient dynamically when valid_sql_rate drops.
+                try:
+                    if hasattr(trainer, "kl_ctl") and hasattr(trainer.kl_ctl, "value"):
+                        trainer.kl_ctl.value *= 1.5
+                        print(f"Increasing KL coef -> {trainer.kl_ctl.value:.4f}")
+                except Exception:
+                    pass
+
+        # ----- logging -----
+        if step % LOG_EVERY == 0:
+            avg_reward = epoch_reward_sum / step
+            print("\n---------------------------")
+            print(f"Epoch {epoch}/{NUM_EPOCHS} | Step {step}/{total_steps} | Avg Reward {avg_reward:.3f}")
+            print("DB:", db_id)
+            print("Q:", question)
+            print("SQL:", pred_sql)
+            print("Reward:", reward_value)
+
+    # epoch stats
+    print(f"\nEpoch {epoch} stats:")
+    print("negative:", negative_rewards)
+    print("partial:", partial_rewards)
+    print("correct:", correct_rewards)
+
+    denom = max(total_considered, 1)
+    print("\nEpoch metrics:")
+    print(f"execution_accuracy: {exec_correct_count/denom:.3f}")
+    print(f"valid_sql_rate: {valid_sql_count/denom:.3f}")
+    print(f"table_match_rate: {table_overlap_sum/denom:.3f}")
+    print(f"column_match_rate: {column_overlap_sum/denom:.3f}")
+    print(f"avg_reward: {epoch_reward_sum/max(denom,1):.3f}")
+    if kl_values:
+        avg_kl = sum(kl_values) / max(len(kl_values), 1)
+        print(f"avg_kl: {avg_kl:.3f}")
+        if avg_kl < -8:
+            print("\nKL collapse guard triggered (avg_kl < -8). Stopping early.")
+            break
+
+    # 🎯 FIXED: Removed the code that saved intermediate checkpoints at the end of each epoch
+
+    # Only save if this epoch is the best one so far
+    epoch_avg_reward = epoch_reward_sum / max(denom, 1)
+    if epoch_avg_reward > best_reward:
+        best_reward = epoch_avg_reward
+        best_epoch = epoch
+
+        print(f"\nNew best model at epoch {epoch} with reward {best_reward:.4f}")
+
+        # 🎯 FIXED: Save directly to checkpoints/rlhf_t5_best, overwriting if needed
+        os.makedirs(output_dir, exist_ok=True)
+
+        trainer.model.save_pretrained(output_dir)
+        tokenizer.save_pretrained(output_dir)
+
+
+print(f"\nTraining finished.")
+print(f"Best epoch: {best_epoch}")
+print(f"Best reward: {best_reward:.4f}")
+print(f"Best model saved at: {output_dir}")

src/train_rl_bart.py

@@ -0,0 +1,370 @@
+# =========================================================
+# RLHF TRAINING FOR TEXT2SQL (OPTIMIZED PPO VERSION - BART)
+# =========================================================
+import torch
+import torch.nn.functional as F
+from datasets import load_dataset
+from transformers import AutoTokenizer
+from trl import PPOTrainer, PPOConfig, AutoModelForSeq2SeqLMWithValueHead
+from peft import PeftModel
+import os, sys, sqlite3, re, random
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+from execution_reward import execution_reward, extract_tables, extract_columns
+
+try:
+    import sqlparse  # gate PPO updates on parsable SQL only
+except Exception:  # pragma: no cover
+    sqlparse = None
+
+# ======================================================
+# DEVICE
+# ======================================================
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
+print("Using device:", device)
+
+# ======================================================
+# TRAINING SETTINGS (🚀 OPTIMIZED FOR SPEED)
+# ======================================================
+NUM_EPOCHS = 10         # Increased to compensate for faster epochs
+LOG_EVERY = 5              # Print logs much more frequently
+MAX_SCHEMA_CHARS = 1500
+MAX_OUTPUT_TOKENS = 48     # 🚀 Down from 64. 95% of Spider SQL is <40 tokens.
+ROLLOUTS_PER_EPOCH = 256   # 🚀 Down from 1024. Epochs will finish 4x faster!
+
+# ======================================================
+# PATHS
+# ======================================================
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DB_ROOT = os.path.join(PROJECT_ROOT, "data/database")
+
+# 🎯 Strict Input: Load strictly from your SFT BART checkpoint
+ADAPTER_PATH = os.path.join(PROJECT_ROOT, "checkpoints/sft_best_bart_2")
+
+# 🎯 Strict Output: Save strictly to rl_best_bart
+OUTPUT_DIR = os.path.join(PROJECT_ROOT, "checkpoints/rl_best_bart")
+
+BASE_MODEL = os.environ.get("BASE_MODEL", "facebook/bart-base")
+
+if not os.path.exists(ADAPTER_PATH):
+    raise RuntimeError(f"❌ No valid LoRA adapter found at: {ADAPTER_PATH}")
+
+print("Loading base:", BASE_MODEL)
+print("Loading adapter:", ADAPTER_PATH)
+
+# ======================================================
+# TOKENIZER
+# ======================================================
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# ======================================================
+# LOAD PPO MODEL
+# ======================================================
+model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float32
+).to(device)
+
+model.pretrained_model = PeftModel.from_pretrained(
+    model.pretrained_model,
+    ADAPTER_PATH,
+    is_trainable=True
+)
+
+# ======================================================
+# LOAD REFERENCE MODEL (FROZEN)
+# ======================================================
+ref_model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float32
+).to(device)
+
+ref_model.pretrained_model = PeftModel.from_pretrained(
+    ref_model.pretrained_model,
+    ADAPTER_PATH,
+    is_trainable=False
+)
+
+ref_model.eval()
+for p in ref_model.parameters():
+    p.requires_grad = False
+
+# ======================================================
+# TRAINABLE PARAMS — ONLY LoRA + VALUE HEAD
+# ======================================================
+for name, p in model.named_parameters():
+    if "lora_" in name or "v_head" in name:
+        p.requires_grad = True
+    else:
+        p.requires_grad = False
+
+model.train()  
+
+trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+total = sum(p.numel() for p in model.parameters())
+print(f"Trainable params: {trainable}/{total} ({100*trainable/total:.2f}%)")
+
+model.config.use_cache = False
+ref_model.config.use_cache = False
+
+# ======================================================
+# DATASET
+# ======================================================
+print("Loading Spider subset...")
+random.seed(0)
+
+TRAIN_DBS = [
+    # already trained
+    "flight_1","student_assessment","store_1","bike_1","book_2","chinook_1",
+    "academic","aircraft","car_1","cinema","club_1","csu_1",
+
+    # medium difficulty (NEW)
+    "college_1","college_2","company_1","company_employee",
+    "customer_complaints","department_store","employee_hire_evaluation",
+    "museum_visit","products_for_hire","restaurant_1",
+    "school_finance","shop_membership","small_bank_1",
+    "soccer_1","student_1","tvshow","voter_1","world_1"
+]
+dataset = load_dataset("spider", split="train")
+dataset = dataset.filter(lambda x: x["db_id"] in TRAIN_DBS)
+
+def valid_example(x):
+    return 5 <= len(x["question"].split()) <= 40
+
+dataset = dataset.filter(valid_example)
+print("Filtered dataset size:", len(dataset))
+
+def sample_example():
+    return dataset[random.randrange(len(dataset))]
+
+# ======================================================
+# DB UTILITIES
+# ======================================================
+def get_db_path(db_id):
+    return os.path.join(DB_ROOT, db_id, f"{db_id}.sqlite")
+
+_SCHEMA_CACHE = {}
+
+def get_db_schema_cached(db_path):
+    if db_path in _SCHEMA_CACHE:
+        return _SCHEMA_CACHE[db_path]
+        
+    schema_text = ""
+    try:
+        conn = sqlite3.connect(db_path)
+        cursor = conn.cursor()
+        tables = cursor.execute("SELECT name FROM sqlite_master WHERE type='table';").fetchall()
+
+        for table in tables:
+            table_name = table[0]
+            columns = cursor.execute(f"PRAGMA table_info({table_name});").fetchall()
+            col_names = [col[1] for col in columns]
+            schema_text += f"{table_name}({', '.join(col_names)})\n"
+        conn.close()
+    except:
+        pass
+        
+    _SCHEMA_CACHE[db_path] = schema_text.strip()
+    return _SCHEMA_CACHE[db_path]
+
+# ======================================================
+# PROMPT
+# ======================================================
+def trim_schema(schema: str, max_chars: int = 1200) -> str:
+    if schema is None:
+        return ""
+    schema = str(schema)
+    if len(schema) <= max_chars:
+        return schema
+    return schema[:max_chars]
+
+def build_prompt(question: str, schema: str) -> str:
+    schema = trim_schema(schema, max_chars=MAX_SCHEMA_CHARS)
+    return f"Database Schema:\n{schema}\n\nTranslate English to SQL:\n{question}\nSQL:\n"
+
+# ======================================================
+# PPO CONFIG (STABLE POLICY LEARNING)
+# ======================================================
+ppo_config = PPOConfig(
+    learning_rate=3e-6,          # slower = prevents policy jump (very important)
+    batch_size=8,
+    mini_batch_size=4,           # good size, keep this
+    gradient_accumulation_steps=2,
+
+    ppo_epochs=2,                # smoother policy update (was 1 → unstable)
+
+    # ---- KL CONTROL (main fix for negative KL) ----
+    init_kl_coef=0.1,
+    target_kl=0.08,              # 0.02 was too strict → caused oscillation
+    adap_kl_ctrl=True,
+
+    # ---- CLIPPING ----
+    cliprange=0.15,
+    cliprange_value=0.15,
+
+    # ---- REWARD STABILITY ----
+    whiten_rewards=True,         # VERY IMPORTANT for binary execution reward
+    kl_penalty="kl",
+
+    # ---- GRADIENT SAFETY ----
+    max_grad_norm=0.3,
+)
+trainer = PPOTrainer(
+    config=ppo_config,
+    model=model,
+    ref_model=ref_model,
+    tokenizer=tokenizer,
+)
+
+try:
+    model.device = torch.device(device)
+except Exception:
+    pass
+
+# ======================================================
+# GENERATION CONFIG
+# ======================================================
+generation_kwargs = dict(
+    max_new_tokens=MAX_OUTPUT_TOKENS,
+    do_sample=True,
+    temperature=0.7,
+    top_p=0.9,
+    pad_token_id=tokenizer.pad_token_id,
+    eos_token_id=tokenizer.eos_token_id,
+)
+# ======================================================
+# TRAIN LOOP (BATCHED & OPTIMIZED)
+# ======================================================
+print("Starting RL training 🚀 (BART PPO Optimized)")
+
+best_reward = -1e9
+global_ppo_step = 0
+model.train()
+
+for epoch in range(1, NUM_EPOCHS + 1):
+    epoch_reward_sum = 0
+    valid_sql_count = 0
+    total_seen = 0
+
+    for step in range(0, ROLLOUTS_PER_EPOCH, ppo_config.batch_size):
+        
+        batch_prompts = []
+        batch_meta = [] 
+
+        for _ in range(ppo_config.batch_size):
+            example = sample_example()
+            question = example["question"]
+            gold_sql = example["query"]
+            db_id = example["db_id"]
+            db_path = get_db_path(db_id)
+
+            schema = get_db_schema_cached(db_path)
+            prompt = build_prompt(question, schema)
+            
+            batch_prompts.append(prompt)
+            batch_meta.append((question, gold_sql, db_path, db_id))
+
+        encoded_inputs = tokenizer(
+            batch_prompts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=512,
+            pad_to_multiple_of=8
+        ).to(device)
+        
+        query_tensors = [encoded_inputs.input_ids[i] for i in range(ppo_config.batch_size)]
+
+        # 🎯 BYPASS: Native model.generate to prevent TRL's truncation crash
+        with torch.no_grad():
+            response_tensors_raw = model.generate(
+                input_ids=encoded_inputs.input_ids,
+                attention_mask=encoded_inputs.attention_mask,
+                **generation_kwargs
+            )
+
+        batch_rewards = []
+        batch_responses_text = []
+        response_tensors = []
+
+        for i in range(ppo_config.batch_size):
+            resp = response_tensors_raw[i]
+            
+            # 🎯 Strip padding safely so TRL's mask calculation never crashes
+            non_pad_mask = resp != tokenizer.pad_token_id
+            if non_pad_mask.sum() == 0:
+                resp = torch.tensor([tokenizer.eos_token_id], device=device)
+                non_pad_mask = resp != tokenizer.pad_token_id
+                
+            valid_len = non_pad_mask.nonzero()[-1].item() + 1
+            clean_resp = resp[:valid_len]
+            response_tensors.append(clean_resp)
+
+            response = tokenizer.decode(clean_resp, skip_special_tokens=True)
+            batch_responses_text.append(response)
+            
+            question, gold_sql, db_path, db_id = batch_meta[i]
+            total_seen += 1
+
+            if "select" not in response.lower():
+                batch_rewards.append(torch.tensor(-1.0, dtype=torch.float32).to(device))
+                continue
+
+            reward = execution_reward(response, db_path, gold_sql)
+            if reward is None:
+                batch_rewards.append(torch.tensor(-1.0, dtype=torch.float32).to(device))
+                continue
+
+            reward = float(reward)
+
+            pred_tables = extract_tables(response)
+            gold_tables = extract_tables(gold_sql)
+            if len(gold_tables) > 0:
+                reward += 0.25 * (len(pred_tables & gold_tables) / len(gold_tables))
+
+            pred_cols = extract_columns(response)
+            gold_cols = extract_columns(gold_sql)
+            if len(gold_cols) > 0:
+                reward += 0.15 * (len(pred_cols & gold_cols) / len(gold_cols))
+
+            reward = max(-1.0, min(1.0, reward))
+            batch_rewards.append(torch.tensor(reward, dtype=torch.float32).to(device))
+            
+            epoch_reward_sum += reward
+            valid_sql_count += 1
+
+        # ---------- PPO UPDATE ----------
+        try:
+            trainer.step(query_tensors, response_tensors, batch_rewards)
+            global_ppo_step += 1
+        except Exception as e:
+            print("⚠️ PPO skipped:", e)
+            continue
+
+        # ---------- LOG ----------
+        if step % (LOG_EVERY * ppo_config.batch_size) == 0 and valid_sql_count > 0:
+            print("\n---------------------------")
+            print(f"Epoch {epoch}/{NUM_EPOCHS} Step {step}/{ROLLOUTS_PER_EPOCH} | Global Update {global_ppo_step}")
+            print("Avg Reward:", round(epoch_reward_sum/valid_sql_count,3))
+            print("Valid SQL:", valid_sql_count,"/",total_seen)
+            
+            sample_idx = random.randint(0, ppo_config.batch_size - 1)
+            print("DB:", batch_meta[sample_idx][3])
+            print("Q:", batch_meta[sample_idx][0])
+            print("SQL:", batch_responses_text[sample_idx])
+            print("Reward:", round(batch_rewards[sample_idx].item(), 3))
+
+    # ---------- SAVE ONLY THE BEST MODEL ----------
+    avg_reward = epoch_reward_sum / max(valid_sql_count, 1)
+
+    if avg_reward > best_reward:
+        best_reward = avg_reward
+        os.makedirs(OUTPUT_DIR, exist_ok=True)
+        
+        model.save_pretrained(OUTPUT_DIR)
+        tokenizer.save_pretrained(OUTPUT_DIR)
+        
+        print(f"\n✅ Saved BEST RLHF model for Epoch {epoch} (reward {best_reward:.3f}) at {OUTPUT_DIR}")

src/train_rl_codet5.py

@@ -0,0 +1,409 @@
+# =========================================================
+# RLHF TRAINING FOR TEXT2SQL (STABLE PPO VERSION)
+# =========================================================
+import torch
+import torch.nn.functional as F
+from datasets import load_dataset
+from transformers import AutoTokenizer
+from transformers.generation.logits_process import LogitsProcessor, LogitsProcessorList
+from trl import PPOTrainer, PPOConfig, AutoModelForSeq2SeqLMWithValueHead
+from peft import PeftModel
+import os, sys, sqlite3, re, random
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+from execution_reward import execution_reward, extract_tables, extract_columns
+try:
+    import sqlparse  # gate PPO updates on parsable SQL only
+except Exception:  # pragma: no cover
+    sqlparse = None
+
+# ======================================================
+# DEVICE
+# ======================================================
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
+print("Using device:", device)
+
+# ======================================================
+# TRAINING SETTINGS
+# ======================================================
+NUM_EPOCHS = 15            
+LOG_EVERY = 20
+USE_SCHEMA = True
+SCHEMA_WARMUP_EPOCHS = 2   
+MAX_SCHEMA_CHARS = 1500
+MAX_OUTPUT_TOKENS = 64     # 🚀 Speed up: Reduced max tokens
+ROLLOUTS_PER_EPOCH = 1024  
+
+# ======================================================
+# PATHS
+# ======================================================
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+RL_MODEL_PATH = os.path.join(PROJECT_ROOT, "outputs/rlhf_text2sql")
+output_dir = RL_MODEL_PATH
+DB_ROOT = os.path.join(PROJECT_ROOT, "data/database")
+
+# Explicit resume checkpoint
+RESUME_CHECKPOINT = os.path.join(PROJECT_ROOT, "checkpoints/milestone_before_more_dbs")
+
+ADAPTER_PATH = os.path.abspath(os.path.join(PROJECT_ROOT, "checkpoints/sft_adapter_codet5"))
+FALLBACK_ADAPTER_PATH = ADAPTER_PATH
+FALLBACK_ADAPTER_PATH_2 = os.path.join(PROJECT_ROOT, "checkpoints")
+
+BASE_MODEL = os.environ.get("BASE_MODEL", "Salesforce/codet5-base")
+
+# ======================================================
+# LOAD MODEL (LoRA)
+# ======================================================
+def find_valid_adapter(path_candidates):
+    # 🚀 SAFETY & RESUME: Check for existing milestone first
+    if os.path.exists(os.path.join(RESUME_CHECKPOINT, "adapter_config.json")):
+        print(f"\n✅ Resuming RL training from checkpoint: {RESUME_CHECKPOINT}\n")
+        return RESUME_CHECKPOINT
+        
+    for p in path_candidates:
+        if p and os.path.exists(os.path.join(p, "adapter_config.json")):
+            return os.path.abspath(p)
+    return None
+
+print("Loading base:", BASE_MODEL)
+
+ADAPTER_PATH = find_valid_adapter([
+    ADAPTER_PATH,
+    FALLBACK_ADAPTER_PATH,
+    FALLBACK_ADAPTER_PATH_2,
+])
+
+if ADAPTER_PATH is None:
+    raise RuntimeError("❌ No valid LoRA adapter found!")
+
+print("Loading adapter:", ADAPTER_PATH)
+
+# ======================================================
+# TOKENIZER
+# ======================================================
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# ======================================================
+# LOAD PPO MODEL
+# ======================================================
+model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float32
+).to(device)
+
+# 🚀 RESUME: Load adapter dynamically and ensure it's trainable
+model.pretrained_model = PeftModel.from_pretrained(
+    model.pretrained_model,
+    ADAPTER_PATH,
+    is_trainable=True
+)
+
+# ======================================================
+# LOAD REFERENCE MODEL (FROZEN)
+# ======================================================
+ref_model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float32
+).to(device)
+
+ref_model.pretrained_model = PeftModel.from_pretrained(
+    ref_model.pretrained_model,
+    ADAPTER_PATH,
+    is_trainable=False
+)
+
+ref_model.eval()
+for p in ref_model.parameters():
+    p.requires_grad = False
+
+# ======================================================
+# TRAINABLE PARAMS — ONLY LoRA + VALUE HEAD
+# ======================================================
+for name, p in model.named_parameters():
+    if "lora_" in name or "v_head" in name:
+        p.requires_grad = True
+    else:
+        p.requires_grad = False
+
+model.train()  
+
+trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+total = sum(p.numel() for p in model.parameters())
+print(f"Trainable params: {trainable}/{total} ({100*trainable/total:.2f}%)")
+
+model.config.use_cache = False
+ref_model.config.use_cache = False
+
+# ======================================================
+# DATASET
+# ======================================================
+print("Loading Spider subset...")
+random.seed(0)
+
+TRAIN_DBS = [
+    # already trained
+    "flight_1","student_assessment","store_1","bike_1","book_2","chinook_1",
+    "academic","aircraft","car_1","cinema","club_1","csu_1",
+
+    # medium difficulty (NEW)
+    "college_1","college_2","company_1","company_employee",
+    "customer_complaints","department_store","employee_hire_evaluation",
+    "museum_visit","products_for_hire","restaurant_1",
+    "school_finance","shop_membership","small_bank_1",
+    "soccer_1","student_1","tvshow","voter_1","world_1"
+]
+
+dataset = load_dataset("spider", split="train")
+dataset = dataset.filter(lambda x: x["db_id"] in TRAIN_DBS)
+
+def valid_example(x):
+    return 5 <= len(x["question"].split()) <= 40
+
+dataset = dataset.filter(valid_example)
+print("Filtered dataset size:", len(dataset))
+
+def sample_example():
+    return dataset[random.randrange(len(dataset))]
+
+# ======================================================
+# DB UTILITIES
+# ======================================================
+def get_db_path(db_id):
+    return os.path.join(DB_ROOT, db_id, f"{db_id}.sqlite")
+
+# 🚀 SPEED OPTIMIZATION: Cache schema so we don't spam disk IO
+_SCHEMA_CACHE = {}
+
+def get_db_schema_cached(db_path):
+    if db_path in _SCHEMA_CACHE:
+        return _SCHEMA_CACHE[db_path]
+        
+    schema_text = ""
+    try:
+        conn = sqlite3.connect(db_path)
+        cursor = conn.cursor()
+        tables = cursor.execute("SELECT name FROM sqlite_master WHERE type='table';").fetchall()
+
+        for table in tables:
+            table_name = table[0]
+            columns = cursor.execute(f"PRAGMA table_info({table_name});").fetchall()
+            col_names = [col[1] for col in columns]
+            schema_text += f"{table_name}({', '.join(col_names)}) "
+        conn.close()
+    except:
+        pass
+        
+    _SCHEMA_CACHE[db_path] = schema_text
+    return schema_text
+
+# ======================================================
+# PROMPT
+# ======================================================
+def trim_schema(schema: str, max_chars: int = 1200) -> str:
+    if schema is None:
+        return ""
+    schema = str(schema)
+    if len(schema) <= max_chars:
+        return schema
+    return schema[:max_chars]
+
+def build_prompt(question: str, schema: str, use_schema: bool) -> str:
+    if not use_schema:
+        return f"### Question:\n{question}\n### SQL:"
+    schema = trim_schema(schema, max_chars=MAX_SCHEMA_CHARS)
+    return f"### Database Schema:\n{schema}\n### Question:\n{question}\n### SQL:"
+
+# ======================================================
+# PPO CONFIG (STABLE POLICY LEARNING)
+# ======================================================
+ppo_config = PPOConfig(
+    learning_rate=5e-6,          
+    batch_size=8,
+    mini_batch_size=2,
+
+    gradient_accumulation_steps=2,
+    ppo_epochs=1,
+    init_kl_coef=0.2,            
+    target_kl=0.02,
+    adap_kl_ctrl=True,
+    cliprange=0.1,
+    cliprange_value=0.1,
+    whiten_rewards=False,
+    kl_penalty="kl",
+    max_grad_norm=0.5,
+)
+
+trainer = PPOTrainer(
+    config=ppo_config,
+    model=model,
+    ref_model=ref_model,
+    tokenizer=tokenizer,
+)
+
+try:
+    model.device = torch.device(device)
+except Exception:
+    pass
+
+# ======================================================
+# GENERATION CONFIG
+# ======================================================
+# 🚀 SPEED OPTIMIZATION: generation limits and randomness bypass
+generation_kwargs = dict(
+    max_new_tokens=MAX_OUTPUT_TOKENS,
+    do_sample=True,          # TRL Requires do_sample=True
+    temperature=1.0,         # Disabled randomness logic 
+    top_p=1.0,               # Disabled randomness logic
+    top_k=0,                 # Disabled randomness logic
+    pad_token_id=tokenizer.pad_token_id,
+    eos_token_id=tokenizer.eos_token_id,
+)
+
+# ======================================================
+# TRAIN LOOP (BATCHED & OPTIMIZED)
+# ======================================================
+print("Starting RL training 🚀 (CodeT5 PPO Stable)")
+
+best_reward = -1e9
+global_ppo_step = 0
+model.train()
+
+for epoch in range(1, NUM_EPOCHS + 1):
+    epoch_reward_sum = 0
+    valid_sql_count = 0
+    total_seen = 0
+
+    # Process in exact chunks matching batch_size to avoid buffer remnants
+    for step in range(0, ROLLOUTS_PER_EPOCH, ppo_config.batch_size):
+        
+        batch_prompts = []
+        batch_meta = [] # Store tuple of (question, gold_sql, db_path, db_id)
+
+        # 🚀 BATCH PREPARATION
+        for _ in range(ppo_config.batch_size):
+            example = sample_example()
+            question = example["question"]
+            gold_sql = example["query"]
+            db_id = example["db_id"]
+            db_path = get_db_path(db_id)
+
+            schema = get_db_schema_cached(db_path)
+            prompt = build_prompt(question, schema, use_schema=True)
+            
+            batch_prompts.append(prompt)
+            batch_meta.append((question, gold_sql, db_path, db_id))
+
+        # 🚀 SPEED OPTIMIZATION: Padded Batch Tokenization (Multiple of 8)
+        encoded_inputs = tokenizer(
+            batch_prompts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=512,
+            pad_to_multiple_of=8
+        ).to(device)
+        
+        # TRL expects lists of 1D tensors
+        query_tensors = [encoded_inputs.input_ids[i] for i in range(ppo_config.batch_size)]
+
+        # 🚀 SPEED OPTIMIZATION: Disable gradients for generation pass
+        with torch.no_grad():
+            response_tensors = trainer.generate(
+                query_tensors,
+                **generation_kwargs
+            )
+
+        batch_rewards = []
+        batch_responses_text = []
+
+        # 🚀 BATCH SQL REWARD EXECUTION (Strictly CPU strings)
+        for i in range(ppo_config.batch_size):
+            response = tokenizer.decode(response_tensors[i], skip_special_tokens=True)
+            batch_responses_text.append(response)
+            question, gold_sql, db_path, db_id = batch_meta[i]
+            
+            total_seen += 1
+
+            # ---------- BASIC SQL FILTER ----------
+            if "select" not in response.lower():
+                batch_rewards.append(torch.tensor(-1.0, dtype=torch.float32).to(device))
+                continue
+
+            # ---------- EXECUTION REWARD ----------
+            reward = execution_reward(response, db_path, gold_sql)
+            if reward is None:
+                batch_rewards.append(torch.tensor(-1.0, dtype=torch.float32).to(device))
+                continue
+
+            reward = float(reward)
+
+            # ---------- TABLE BONUS ----------
+            pred_tables = extract_tables(response)
+            gold_tables = extract_tables(gold_sql)
+            if len(gold_tables) > 0:
+                reward += 0.25 * (len(pred_tables & gold_tables) / len(gold_tables))
+
+            # ---------- COLUMN BONUS ----------
+            pred_cols = extract_columns(response)
+            gold_cols = extract_columns(gold_sql)
+            if len(gold_cols) > 0:
+                reward += 0.15 * (len(pred_cols & gold_cols) / len(gold_cols))
+
+            # ---------- CLAMP ----------
+            reward = max(-1.0, min(1.0, reward))
+            batch_rewards.append(torch.tensor(reward, dtype=torch.float32).to(device))
+            
+            epoch_reward_sum += reward
+            valid_sql_count += 1
+
+        # ---------- PPO UPDATE ----------
+        try:
+            trainer.step(
+                query_tensors,
+                response_tensors,
+                batch_rewards
+            )
+            global_ppo_step += 1
+        except Exception as e:
+            print("⚠️ PPO skipped:", e)
+            continue
+
+        # 🚀 AUTO CHECKPOINT SAVING: Every 200 PPO Updates
+        if global_ppo_step > 0 and global_ppo_step % 200 == 0:
+            step_save_path = os.path.join(PROJECT_ROOT, f"checkpoints/rl_step_{global_ppo_step}")
+            os.makedirs(step_save_path, exist_ok=True)
+            
+            # Saves ONLY the adapter, keeping disk usage tiny!
+            model.save_pretrained(step_save_path)
+            tokenizer.save_pretrained(step_save_path)
+            print(f"\n💾 [AUTO-SAVE] Checkpoint saved at PPO step {global_ppo_step} -> {step_save_path}")
+
+        # ---------- LOG ----------
+        if step % (LOG_EVERY * ppo_config.batch_size) == 0 and valid_sql_count > 0:
+            print("\n---------------------------")
+            print(f"Epoch {epoch}/{NUM_EPOCHS} Step {step}/{ROLLOUTS_PER_EPOCH} | Global Update {global_ppo_step}")
+            print("Avg Reward:", round(epoch_reward_sum/valid_sql_count,3))
+            print("Valid SQL:", valid_sql_count,"/",total_seen)
+            
+            # Print sample from latest batch
+            sample_idx = random.randint(0, ppo_config.batch_size - 1)
+            print("DB:", batch_meta[sample_idx][3])
+            print("Q:", batch_meta[sample_idx][0])
+            print("SQL:", batch_responses_text[sample_idx])
+            print("Reward:", round(batch_rewards[sample_idx].item(), 3))
+
+    # ---------- SAVE BEST MODEL (INSIDE EPOCH) ----------
+    avg_reward = epoch_reward_sum / max(valid_sql_count, 1)
+
+    if avg_reward > best_reward:
+        best_reward = avg_reward
+        save_path = os.path.join(PROJECT_ROOT, "checkpoints/best_rlhf_model")
+        os.makedirs(save_path, exist_ok=True)
+        
+        model.save_pretrained(save_path)
+        tokenizer.save_pretrained(save_path)
+        
+        print(f"\n✅ Saved BEST RLHF model for Epoch {epoch} (reward {best_reward:.3f})")

src/train_rl_lora.py

@@ -0,0 +1,151 @@
+# ======================================
+# RLHF Text2SQL — FINAL WORKING VERSION
+# T5-small + LoRA + PPO + Execution Reward
+# Single-sample stable training (Mac MPS safe)
+# ======================================
+
+from execution_reward import execution_reward
+import os, gc, json, random, torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from trl import PPOTrainer, PPOConfig
+from trl.models.modeling_value_head import AutoModelForSeq2SeqLMWithValueHead
+from peft import LoraConfig, get_peft_model
+
+# ---------------- SETTINGS ----------------
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+device = "mps" if torch.backends.mps.is_available() else "cpu"
+print("Using device:", device)
+
+os.makedirs("rlhf_text2sql_lora", exist_ok=True)
+
+# ---------------- MODEL ----------------
+model_name = "google/flan-t5-small"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer.pad_token = tokenizer.eos_token
+
+base_model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+# LoRA
+lora_config = LoraConfig(
+    r=8,
+    lora_alpha=16,
+    target_modules=["q","v"],
+    lora_dropout=0.05,
+    bias="none",
+    task_type="SEQ_2_SEQ_LM",
+)
+
+base_model = get_peft_model(base_model, lora_config)
+
+model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(base_model).to(device)
+ref_model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(model_name).to(device)
+
+model.config.use_cache = False
+ref_model.config.use_cache = False
+
+# ---------------- DATA ----------------
+with open("data/train_spider.json") as f:
+    dataset = json.load(f)
+
+def build_prompt(example):
+    return f"Translate to SQL: {example['question']}"
+
+# ---------------- PPO ----------------
+ppo_config = PPOConfig(
+    batch_size=1,
+    mini_batch_size=1,
+    learning_rate=2e-6,
+    target_kl=0.05,
+    adap_kl_ctrl=True,
+    init_kl_coef=0.2,
+)
+
+ppo_trainer = PPOTrainer(
+    config=ppo_config,
+    model=model,
+    ref_model=ref_model,
+    tokenizer=tokenizer,
+)
+
+# ---------------- GENERATION ----------------
+def generate_sql(query_tensors):
+
+    # deterministic decoding = prevents NaN explosion
+    with torch.no_grad():
+        response_tensors = ppo_trainer.generate(
+            query_tensors,
+            max_new_tokens=64,
+
+            # 🔴 CRITICAL: disable sampling
+            do_sample=False,
+
+            # stable decoding
+            num_beams=1,
+            early_stopping=True,
+
+            # prevents invalid tokens
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    # extra safety (important on MPS)
+    cleaned = []
+    for t in response_tensors:
+        t = torch.nan_to_num(t, nan=0, posinf=0, neginf=0)
+        cleaned.append(t)
+
+    return cleaned
+
+# ---------------- TRAIN ----------------
+MAX_STEPS = 1200
+
+for step in range(MAX_STEPS):
+
+    # pick random Spider example
+    example = random.choice(dataset)
+
+    question = example["question"]
+    gold_sql = example["query"]
+    db_id = example["db_id"]
+    db_path = f"data/database/{db_id}/{db_id}.sqlite"
+
+    # tokenize
+    enc = tokenizer(build_prompt(example), return_tensors="pt")
+    query_tensor = enc.input_ids.to(device)
+    query_tensors = [query_tensor[0]]
+
+    # generate SQL
+    response_tensors = generate_sql(query_tensors)
+    pred_sql = tokenizer.decode(response_tensors[0], skip_special_tokens=True)
+
+    # -------- EXECUTION REWARD --------
+    reward = execution_reward(pred_sql, gold_sql, db_path)
+    reward_tensor = torch.tensor([reward], dtype=torch.float32).to(device)
+
+    # PPO update
+    stats = ppo_trainer.step(query_tensors, response_tensors, [reward_tensor])
+
+    # stabilize
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+
+    # cleanup
+    del query_tensor, response_tensors, reward_tensor
+    gc.collect()
+    if device == "mps":
+        torch.mps.empty_cache()
+
+    # log
+    if step % 20 == 0:
+        print(f"\nStep {step}/{MAX_STEPS}")
+        print("DB:", db_id)
+        print("Q:", question)
+        print("Pred:", pred_sql)
+        print("Gold:", gold_sql)
+        print("Reward:", reward)
+
+# ---------------- SAVE ----------------
+model.save_pretrained("rlhf_text2sql_lora")
+tokenizer.save_pretrained("rlhf_text2sql_lora")
+
+print("\nTraining complete — model saved!")

src/train_sft.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+import os
+import torch
+from datasets import load_dataset
+from peft import LoraConfig, get_peft_model
+from transformers import (
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    DataCollatorForSeq2Seq,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+)
+
+from prompting import clean_gold_sql, get_schema_text, build_prompt
+
+# =====================================================
+# SETTINGS
+# =====================================================
+BASE_MODEL = os.environ.get("BASE_MODEL", "t5-small")
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+# 🎯 FIXED: Save final model to checkpoints/sft_t5 to protect existing models
+OUT_DIR = os.path.join(PROJECT_ROOT, "checkpoints", "sft_t5")
+
+TRAIN_SPLIT = "train[:7000]"
+EPOCHS = 8
+LR = 3e-4
+PER_DEVICE_BATCH = 4
+GRAD_ACCUM = 2
+
+MAX_INPUT = 512
+MAX_OUTPUT = 128
+
+# =====================================================
+# DEVICE
+# =====================================================
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+print("Using device:", device)
+
+# =====================================================
+# TOKENIZER
+# =====================================================
+print("Loading tokenizer/model:", BASE_MODEL)
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# =====================================================
+# PREPROCESS FUNCTION (CRITICAL FIXED VERSION)
+# =====================================================
+def preprocess_function(example):
+
+    question = example["question"]
+    db_id = example["db_id"]
+    gold_sql = clean_gold_sql(example["query"])
+
+    # ---- Build Prompt ----
+    schema_text = get_schema_text(db_id)
+    prompt = build_prompt(question, db_id, schema_text=schema_text, training_sql=None)
+
+    model_inputs = tokenizer(
+        prompt,
+        max_length=MAX_INPUT,
+        truncation=True,
+        padding="max_length",
+    )
+
+    # ---- Target SQL ----
+    labels = tokenizer(
+        gold_sql,
+        max_length=MAX_OUTPUT,
+        truncation=True,
+        padding="max_length",
+    )["input_ids"]
+
+    # IMPORTANT: ignore padding in loss
+    labels = [
+        (tok if tok != tokenizer.pad_token_id else -100)
+        for tok in labels
+    ]
+
+    model_inputs["labels"] = labels
+    return model_inputs
+
+# =====================================================
+# DATASET
+# =====================================================
+print("Loading Spider subset:", TRAIN_SPLIT)
+dataset = load_dataset("spider", split=TRAIN_SPLIT)
+dataset = dataset.train_test_split(test_size=0.1, seed=42)
+
+train_ds = dataset["train"]
+eval_ds = dataset["test"]
+
+print("Tokenizing dataset (single process, stable)...")
+
+train_tok = train_ds.map(
+    preprocess_function,
+    batched=False,
+    num_proc=1,                    # 🔥 VERY IMPORTANT FIX
+    remove_columns=train_ds.column_names,
+    load_from_cache_file=False,
+)
+
+eval_tok = eval_ds.map(
+    preprocess_function,
+    batched=False,
+    num_proc=1,
+    remove_columns=eval_ds.column_names,
+    load_from_cache_file=False,
+)
+
+print("Train dataset size:", len(train_tok))
+print("Eval dataset size:", len(eval_tok))
+
+# =====================================================
+# MODEL + LoRA
+# =====================================================
+base_model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)
+
+base_model.config.use_cache = False
+base_model.gradient_checkpointing_enable()
+
+lora_config = LoraConfig(
+    r=8,
+    lora_alpha=16,
+    lora_dropout=0.1,
+    bias="none",
+    task_type="SEQ_2_SEQ_LM",
+    target_modules=["q", "v"],   # correct for T5
+)
+
+model = get_peft_model(base_model, lora_config)
+model.to(device)
+
+# =====================================================
+# TRAINER
+# =====================================================
+data_collator = DataCollatorForSeq2Seq(
+    tokenizer=tokenizer,
+    model=model,
+    padding=True,
+)
+
+args = Seq2SeqTrainingArguments(
+    # 🎯 FIXED: Changed path to prevent mixing logs with your old CodeT5 logs
+    output_dir=os.path.join(PROJECT_ROOT, "checkpoints", "sft_t5_runs"),
+    num_train_epochs=EPOCHS,
+    learning_rate=LR,
+    per_device_train_batch_size=PER_DEVICE_BATCH,
+    per_device_eval_batch_size=PER_DEVICE_BATCH,
+    gradient_accumulation_steps=GRAD_ACCUM,
+    dataloader_num_workers=0,
+    dataloader_pin_memory=False,
+    evaluation_strategy="epoch",
+    
+    # 🎯 FIXED: "no" completely stops intermediate saving! Only the final model will be saved.
+    save_strategy="no", 
+    
+    logging_steps=50,
+    report_to=[],
+    fp16=False,
+    bf16=False,
+    predict_with_generate=True,
+)
+
+trainer = Seq2SeqTrainer(
+    model=model,
+    args=args,
+    train_dataset=train_tok,
+    eval_dataset=eval_tok,
+    tokenizer=tokenizer,
+    data_collator=data_collator,
+)
+
+# =====================================================
+# TRAIN
+# =====================================================
+trainer.train()
+
+# =====================================================
+# SAVE
+# =====================================================
+print("Saving LoRA adapter to:", OUT_DIR)
+os.makedirs(OUT_DIR, exist_ok=True)
+model.save_pretrained(OUT_DIR)
+tokenizer.save_pretrained(OUT_DIR)
+
+print("DONE ✔ SFT warmup finished")

src/train_sft_bart.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+import os
+import torch
+from datasets import load_dataset
+from peft import LoraConfig, get_peft_model
+from transformers import (
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    DataCollatorForSeq2Seq,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+)
+
+from prompting import clean_gold_sql, get_schema_text, build_prompt
+
+# =====================================================
+# SETTINGS
+# =====================================================
+BASE_MODEL = os.environ.get("BASE_MODEL", "facebook/bart-base")
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+OUT_DIR = os.path.join(PROJECT_ROOT, "checkpoints", "sft_best_bart_2")
+
+TRAIN_SPLIT = "train[:7000]"
+
+EPOCHS = 12
+LR = 3e-4
+PER_DEVICE_BATCH = 16
+GRAD_ACCUM = 4
+
+MAX_INPUT = 512
+MAX_OUTPUT = 128
+
+# =====================================================
+# DEVICE
+# =====================================================
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+device = torch.device("mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu"))
+print("Using device:", device)
+
+# =====================================================
+# TOKENIZER
+# =====================================================
+print("Loading tokenizer/model:", BASE_MODEL)
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# =====================================================
+# PREPROCESS FUNCTION
+# =====================================================
+def preprocess_function(example):
+
+    question = example["question"]
+    db_id = example["db_id"]
+    gold_sql = clean_gold_sql(example["query"])
+
+    # ---- Build Prompt ----
+    schema_text = get_schema_text(db_id)
+    prompt = build_prompt(question, db_id, schema_text=schema_text, training_sql=None)
+
+    model_inputs = tokenizer(
+        prompt,
+        max_length=MAX_INPUT,
+        truncation=True,
+        padding="max_length",
+    )
+
+    # ---- Target SQL ----
+    labels = tokenizer(
+        gold_sql,
+        max_length=MAX_OUTPUT,
+        truncation=True,
+        padding="max_length",
+    )["input_ids"]
+
+    # IMPORTANT: ignore padding in loss
+    labels = [
+        (tok if tok != tokenizer.pad_token_id else -400)
+        for tok in labels
+    ]
+
+    model_inputs["labels"] = labels
+    return model_inputs
+
+# =====================================================
+# DATASET
+# =====================================================
+print("Loading Spider subset:", TRAIN_SPLIT)
+dataset = load_dataset("spider", split=TRAIN_SPLIT)
+dataset = dataset.train_test_split(test_size=0.1, seed=42)
+
+train_ds = dataset["train"]
+eval_ds = dataset["test"]
+
+print("Tokenizing dataset (single process, stable)...")
+
+train_tok = train_ds.map(
+    preprocess_function,
+    batched=False,
+    num_proc=1,
+    remove_columns=train_ds.column_names,
+    load_from_cache_file=False,
+)
+
+eval_tok = eval_ds.map(
+    preprocess_function,
+    batched=False,
+    num_proc=1,
+    remove_columns=eval_ds.column_names,
+    load_from_cache_file=False,
+)
+
+print("Train dataset size:", len(train_tok))
+print("Eval dataset size:", len(eval_tok))
+
+# =====================================================
+# MODEL + LoRA
+# =====================================================
+base_model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)
+
+base_model.config.use_cache = False
+
+# 🚀 UPGRADE 1: Expanded LoRA brainpower
+lora_config = LoraConfig(
+    r=16,            # Increased rank for more learning capacity
+    lora_alpha=32,   # Alpha is typically 2x the rank
+    lora_dropout=0.1,
+    bias="none",
+    task_type="SEQ_2_SEQ_LM",
+    # Target all attention and dense layers in BART
+    target_modules=["q_proj", "k_proj", "v_proj", "out_proj", "fc1", "fc2"],   
+)
+
+model = get_peft_model(base_model, lora_config)
+model.to(device)
+
+# =====================================================
+# TRAINER
+# =====================================================
+data_collator = DataCollatorForSeq2Seq(
+    tokenizer=tokenizer,
+    model=model,
+    padding=True,
+)
+
+args = Seq2SeqTrainingArguments(
+    output_dir=os.path.join(PROJECT_ROOT, "checkpoints", "sft_bart_runs"),
+    num_train_epochs=EPOCHS,
+    learning_rate=LR,
+    per_device_train_batch_size=PER_DEVICE_BATCH,
+    per_device_eval_batch_size=PER_DEVICE_BATCH,
+    gradient_accumulation_steps=GRAD_ACCUM,
+    dataloader_num_workers=0,
+    dataloader_pin_memory=False,
+    
+    # 🚀 UPGRADE 2 & 3: Better optimization & generalization
+    warmup_ratio=0.05,              # Slowly ramp up learning rate
+    weight_decay=0.01,              # Penalize over-reliance on single tokens
+    label_smoothing_factor=0.1,     # Prevent overconfidence in SQL token matching
+    
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    
+    save_total_limit=1,
+    load_best_model_at_end=True,
+    metric_for_best_model="eval_loss",
+    greater_is_better=False, 
+    
+    logging_steps=50,
+    report_to=[],
+    fp16=False,
+    bf16=False,
+    predict_with_generate=True,
+)
+
+trainer = Seq2SeqTrainer(
+    model=model,
+    args=args,
+    train_dataset=train_tok,
+    eval_dataset=eval_tok,
+    tokenizer=tokenizer,
+    data_collator=data_collator,
+)
+
+# =====================================================
+# TRAIN
+# =====================================================
+trainer.train()
+
+# =====================================================
+# SAVE BEST MODEL
+# =====================================================
+print("Saving best BART LoRA adapter to:", OUT_DIR)
+os.makedirs(OUT_DIR, exist_ok=True)
+
+trainer.model.save_pretrained(OUT_DIR)
+tokenizer.save_pretrained(OUT_DIR)
+
+print("DONE ✔ SFT BART finished")

src/train_sft_codet5.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import os
+import torch
+from datasets import load_dataset
+from peft import LoraConfig, get_peft_model
+from transformers import (
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    DataCollatorForSeq2Seq,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+)
+
+from prompting import clean_gold_sql, get_schema_text, build_prompt
+
+# =====================================================
+# SETTINGS
+# =====================================================
+BASE_MODEL = "Salesforce/codet5-base"
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+OUT_DIR = os.path.join(PROJECT_ROOT, "checkpoints", "sft_adapter_codet5")
+
+TRAIN_SPLIT = "train[:7000]"
+EPOCHS = 10
+LR = 2e-4
+PER_DEVICE_BATCH = 2        # codet5 bigger -> reduce
+GRAD_ACCUM = 4
+
+MAX_INPUT = 512
+MAX_OUTPUT = 160
+
+# =====================================================
+# DEVICE
+# =====================================================
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+print("Using device:", device)
+
+# =====================================================
+# TOKENIZER
+# =====================================================
+print("Loading tokenizer/model:", BASE_MODEL)
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# =====================================================
+# PREPROCESS FUNCTION
+# =====================================================
+def preprocess_function(example):
+
+    question = example["question"]
+    db_id = example["db_id"]
+    gold_sql = clean_gold_sql(example["query"])
+
+    schema_text = get_schema_text(db_id)
+    prompt = build_prompt(question, db_id, schema_text=schema_text, training_sql=None)
+
+    model_inputs = tokenizer(
+        prompt,
+        max_length=MAX_INPUT,
+        truncation=True,
+        padding="max_length",
+    )
+
+    labels = tokenizer(
+        gold_sql,
+        max_length=MAX_OUTPUT,
+        truncation=True,
+        padding="max_length",
+    )["input_ids"]
+
+    labels = [(tok if tok != tokenizer.pad_token_id else -100) for tok in labels]
+
+    model_inputs["labels"] = labels
+    return model_inputs
+
+# =====================================================
+# DATASET
+# =====================================================
+print("Loading Spider subset:", TRAIN_SPLIT)
+dataset = load_dataset("spider", split=TRAIN_SPLIT)
+dataset = dataset.train_test_split(test_size=0.1, seed=42)
+
+train_ds = dataset["train"]
+eval_ds = dataset["test"]
+
+print("Tokenizing dataset...")
+
+train_tok = train_ds.map(
+    preprocess_function,
+    batched=False,
+    num_proc=1,
+    remove_columns=train_ds.column_names,
+    load_from_cache_file=False,
+)
+
+eval_tok = eval_ds.map(
+    preprocess_function,
+    batched=False,
+    num_proc=1,
+    remove_columns=eval_ds.column_names,
+    load_from_cache_file=False,
+)
+
+print("Train dataset size:", len(train_tok))
+print("Eval dataset size:", len(eval_tok))
+
+# =====================================================
+# MODEL + LoRA (CODET5 FIXED)
+# =====================================================
+base_model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)
+
+base_model.config.use_cache = False
+base_model.gradient_checkpointing_enable()
+
+# 🔥 DIFFERENT FROM T5
+lora_config = LoraConfig(
+    r=16,
+    lora_alpha=32,
+    lora_dropout=0.05,
+    bias="none",
+    task_type="SEQ_2_SEQ_LM",
+    target_modules=["q", "v"],   # IMPORTANT FOR CODET5
+)
+
+model = get_peft_model(base_model, lora_config)
+model.to(device)
+
+model.print_trainable_parameters()
+
+# =====================================================
+# TRAINER
+# =====================================================
+data_collator = DataCollatorForSeq2Seq(
+    tokenizer=tokenizer,
+    model=model,
+    padding=True,
+)
+
+args = Seq2SeqTrainingArguments(
+    output_dir=os.path.join(PROJECT_ROOT, "checkpoints", "sft_runs_codet5"),
+    num_train_epochs=EPOCHS,
+    learning_rate=LR,
+    per_device_train_batch_size=PER_DEVICE_BATCH,
+    per_device_eval_batch_size=PER_DEVICE_BATCH,
+    gradient_accumulation_steps=GRAD_ACCUM,
+    dataloader_num_workers=0,
+    dataloader_pin_memory=False,
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    save_total_limit=1,
+    logging_steps=50,
+    report_to=[],
+    fp16=False,
+    bf16=False,
+    predict_with_generate=True,
+)
+
+trainer = Seq2SeqTrainer(
+    model=model,
+    args=args,
+    train_dataset=train_tok,
+    eval_dataset=eval_tok,
+    tokenizer=tokenizer,
+    data_collator=data_collator,
+)
+
+# =====================================================
+# TRAIN
+# =====================================================
+trainer.train()
+
+# =====================================================
+# SAVE
+# =====================================================
+# =====================================================
+# SAVE (SAFE PEFT SAVE)
+# =====================================================
+print("Saving LoRA adapter to:", OUT_DIR)
+os.makedirs(OUT_DIR, exist_ok=True)
+
+# unwrap trainer model (important!)
+peft_model = trainer.model
+
+# ensure on cpu before saving (mac mps bug fix)
+peft_model = peft_model.to("cpu")
+
+# save adapter only
+peft_model.save_pretrained(OUT_DIR)
+tokenizer.save_pretrained(OUT_DIR)
+
+print("DONE ✔ CodeT5 SFT finished")