Add dataset files using LFS

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+llmsql_2.0/*.jsonl filter=lfs diff=lfs merge=lfs -text

app.py

@@ -14,7 +14,16 @@ from build_prompt import (
     build_prompt_1shot,
     build_prompt_5shot,
 )
-from dataset_generator import questions, split_info, tables
+
+# Import both versions of the dataset
+from dataset_generator import (
+    questions_llmsql_1,
+    questions_llmsql_2,
+    tables_llmsql_1,
+    tables_llmsql_2,
+)
+from dataset_generator import split_info as split_info_v1
+from dataset_generator import split_info as split_info_v2
 from evaluate import evaluate_sample
 
 # Global variables for caching
@@ -22,28 +31,56 @@ model = None
 tokenizer = None
 current_model_id = None
 conn = None
+current_db_path = None
+
+# Mapping for dynamic access
+DATASETS = {
+    "LLMSQL 1.0": {
+        "questions": questions_llmsql_1,
+        "tables": tables_llmsql_1,
+        "split_info": split_info_v1,
+        "repo": "llmsql-bench/llmsql-benchmark",
+        "folder": "llmsql_1.0",
+    },
+    "LLMSQL 2.0": {
+        "questions": questions_llmsql_2,
+        "tables": tables_llmsql_2,
+        "split_info": split_info_v2,
+        "repo": "llmsql-bench/llmsql-2.0",
+        "folder": "llmsql_2.0",
+    },
+}
 
 
 # =====================
 # Initialization Logic
 # =====================
-def initialize_data():
-    """Loads datasets once and calculates actual key ranges."""
-    global conn
+def initialize_data(version):
+    """Downloads and connects to the version-specific database."""
+    global conn, current_db_path
 
+    config = DATASETS[version]
     APP_DIR = os.getcwd()
-    DB_FILE = os.path.join(APP_DIR, "sqlite_tables.db")
-
-    if not os.path.exists(DB_FILE):
-        print("Downloading database...")
-        hf_hub_download(
-            repo_id="llmsql-bench/llmsql-benchmark",
-            repo_type="dataset",
-            filename="sqlite_tables.db",
-            local_dir=APP_DIR,
-        )
+    VER_DIR = os.path.join(APP_DIR, config["folder"])
+    os.makedirs(VER_DIR, exist_ok=True)
+
+    DB_FILE = os.path.join(VER_DIR, "sqlite_tables.db")
+
+    # Only reconnect if the version changed or conn is None
+    if current_db_path != DB_FILE:
+        if not os.path.exists(DB_FILE):
+            print(f"Downloading database for {version}...")
+            hf_hub_download(
+                repo_id=config["repo"],
+                repo_type="dataset",
+                filename="sqlite_tables.db",
+                local_dir=VER_DIR,
+            )
 
-    conn = sqlite3.connect(DB_FILE, check_same_thread=False)
+        if conn:
+            conn.close()
+        conn = sqlite3.connect(DB_FILE, check_same_thread=False)
+        current_db_path = DB_FILE
 
 
 def load_model_if_needed(model_id):
@@ -65,9 +102,6 @@ def load_model_if_needed(model_id):
     print(f"Model {model_id} loaded successfully.")
 
 
-# Pre-load data once at startup (CPU tasks)
-initialize_data()
-
 few_shot_selection = {
     "0": build_prompt_0shot,
     "1": build_prompt_1shot,
@@ -79,53 +113,47 @@ few_shot_selection = {
 # Main Logic
 # =====================
 @spaces.GPU
-def run_inference(model_id, question_idx, few_shots):
-    # Ensure data is ready and model is correct one
-    initialize_data()
+def run_inference(version, model_id, question_idx, few_shots):
+    initialize_data(version)
     load_model_if_needed(model_id)
 
+    dataset = DATASETS[version]
+    qs = dataset["questions"]
+    ts = dataset["tables"]
+
     try:
-        question_idx = int(question_idx)
+        idx = int(question_idx)
+        q_data = qs[idx]
     except:
-        return "Question id should be integer.", "", "", None, None, None
+        return "Invalid ID", "", "", None, None, None, False
 
-    # Extraction
-    question = questions[question_idx]["question"]
-    ground_truth_sql = questions[question_idx]["sql"]
-    table_id = questions[question_idx]["table_id"]
-    table = tables.get(table_id)
+    question = q_data["question"]
+    ground_truth_sql = q_data["sql"]
+    table = ts.get(q_data["table_id"])
 
     if not table:
-        return "Table data missing", "", "", None, None, None
+        return "Table data missing", "", "", None, None, None, False
 
     example_row = table["rows"][0] if table["rows"] else []
-
-    # Prompt Engineering
     raw_prompt = few_shot_selection[few_shots](
         question, table["header"], table["types"], example_row
     )
 
-    messages = [
-        {"role": "user", "content": raw_prompt},
-    ]
-
+    messages = [{"role": "user", "content": raw_prompt}]
     text_input = tokenizer.apply_chat_template(
         messages, tokenize=False, add_generation_prompt=True
     )
     model_inputs = tokenizer([text_input], return_tensors="pt").to("cuda")
 
-    # Generation
     generated_ids = model.generate(
         **model_inputs, max_new_tokens=512, temperature=0.0, do_sample=False
     )
-
-    # Slice off the prompt tokens
     new_tokens = generated_ids[0][len(model_inputs.input_ids[0]) :]
     completion = tokenizer.decode(new_tokens, skip_special_tokens=True)
 
-    is_match, mismatch_info, error_info = evaluate_sample(
-        item={"question_id": question_idx, "completion": completion},
-        questions=questions,
+    is_match, mismatch_info, _ = evaluate_sample(
+        item={"question_id": idx, "completion": completion},
+        questions=qs,
         conn=conn,
     )
 
@@ -143,32 +171,31 @@ def run_inference(model_id, question_idx, few_shots):
 # =====================
 # UI Helpers
 # =====================
-def get_range_text(split_name):
-    info = split_info.get(split_name, {})
-    return (
-        f"IDs: {info.get('first')} to {info.get('last')} (Total: {info.get('count')})"
-    )
+def get_range_display(version):
+    info_dict = DATASETS[version]["split_info"]
+    lines = []
+    for s, info in info_dict.items():
+        lines.append(
+            f"**{s.capitalize()}**: IDs {info.get('first')} to {info.get('last')} (Total: {info.get('count')})"
+        )
+    return "\n".join(lines)
 
 
-# Helper to build your static range text
-range_display = "\n".join(
-    [f"**{s.capitalize()}**: {get_range_text(s)}" for s in split_info.keys()]
-)
-
 with gr.Blocks(title="Text-to-SQL Debugger", theme=gr.themes.Soft()) as app:
     gr.Markdown("## 🔍 Text-to-SQL Interactive Debugger")
 
     with gr.Row():
+        version_dropdown = gr.Dropdown(
+            choices=["LLMSQL 1.0", "LLMSQL 2.0"],
+            value="LLMSQL 2.0",
+            label="Dataset Version",
+            scale=1,
+        )
         model_dropdown = gr.Dropdown(
-            choices=[
-                "Qwen/Qwen2.5-1.5B-Instruct",
-                # "microsoft/Phi-3.5-mini-instruct",
-                "openai/gpt-oss-20b",
-                # "Qwen/Qwen3-Coder-30B-A3B-Instruct",
-            ],
+            choices=["Qwen/Qwen2.5-1.5B-Instruct", "openai/gpt-oss-20b"],
             value="Qwen/Qwen2.5-1.5B-Instruct",
             label="1. Select Model",
-            scale=2,
+            scale=1,
         )
         few_shot_dropdown = gr.Dropdown(
             choices=["0", "1", "5"],
@@ -182,20 +209,12 @@ with gr.Blocks(title="Text-to-SQL Debugger", theme=gr.themes.Soft()) as app:
             label="3. Enter Question ID",
             value="1",
             lines=2,
-            placeholder="e.g. 15001",
+            placeholder="e.g. 1",
             scale=2,
-            min_width=200,
         )
-        with gr.Column(scale=1):
-            gr.Markdown(
-                f"**Available ID Ranges:**\n{range_display}",
-                line_breaks=True,
-                padding=True,
-            )
-
-    with gr.Row():
-        run_button = gr.Button("Run Inference", variant="primary")
+        range_md = gr.Markdown(get_range_display("LLMSQL 2.0"), scale=1)
 
+    run_button = gr.Button("Run Inference", variant="primary")
     question_box = gr.Textbox(
         label="Natural Language Question", lines=2, interactive=False
     )
@@ -206,69 +225,76 @@ with gr.Blocks(title="Text-to-SQL Debugger", theme=gr.themes.Soft()) as app:
 
     gr.Markdown("### Data Comparison")
     with gr.Row():
-        generated_table = gr.Dataframe(
-            label="Generated Result", type="pandas", wrap=True
-        )
-        gt_table = gr.Dataframe(label="Ground Truth Result", type="pandas", wrap=True)
+        generated_table = gr.Dataframe(label="Generated Result", type="pandas")
+        gt_table = gr.Dataframe(label="Ground Truth Result", type="pandas")
 
     with gr.Accordion("See Full Source Table", open=False):
         full_table = gr.Dataframe(label="Full Table Content", type="pandas")
 
-    def update_question_info(q_id):
+    def update_ui_on_version_or_id(version, q_id):
+        """Updates range text and pre-loads question data when version or ID changes."""
+        dataset = DATASETS[version]
+        range_text = get_range_display(version)
+
         try:
             idx = int(q_id) if (q_id and str(q_id).isdigit()) else 1
-            question = questions[idx]["question"]
-            gt_sql = questions[idx]["sql"]
-            table_id = questions[idx]["table_id"]
-
-            raw_table_data = tables.get(table_id)
+            q_data = dataset["questions"][idx]
+            table_id = q_data["table_id"]
+            raw_table = dataset["tables"].get(table_id, {})
 
-            if raw_table_data and "rows" in raw_table_data:
-                df = pd.DataFrame(
-                    raw_table_data["rows"], columns=raw_table_data["header"]
-                )
-            else:
-                df = pd.DataFrame()
-
-            return (question, gt_sql, df, gr.update(label="Generated SQL", value=""))
-        except Exception as e:
-            print(f"Error: {e}")  # Helpful for debugging
-            return "ID not found", "", None, gr.update(label="Generated SQL")
+            df = pd.DataFrame(
+                raw_table.get("rows", []), columns=raw_table.get("header", [])
+            )
+            return (
+                q_data["question"],
+                q_data["sql"],
+                df,
+                gr.update(label="Generated SQL", value=""),
+                range_text,
+            )
+        except Exception:
+            return (
+                "ID not found in this version",
+                "",
+                pd.DataFrame(),
+                gr.update(label="Generated SQL"),
+                range_text,
+            )
 
-    def handle_inference(model, few_shot, q_id):
+    def handle_inference(version, model, few_shot, q_id):
         q_text, gen_sql, gt_sql, gen_df, gt_df, full_df, is_match = run_inference(
-            model, q_id, few_shot
+            version, model, q_id, few_shot
         )
-
-        status_label = (
+        status = (
             "✅ Generated SQL (MATCH SUCCESS)"
             if is_match
             else "❌ Generated SQL (MATCH FAILED)"
         )
+        return gr.update(label=status, value=gen_sql), gen_df, gt_df
 
-        return (gr.update(label=status_label, value=gen_sql), gen_df, gt_df)
+    # Event Listeners
+    version_dropdown.change(
+        update_ui_on_version_or_id,
+        inputs=[version_dropdown, question_input],
+        outputs=[question_box, gt_sql_box, full_table, generated_sql_box, range_md],
+    )
 
     question_input.change(
-        fn=update_question_info,
-        inputs=[question_input],
-        outputs=[question_box, gt_sql_box, full_table, generated_sql_box],
+        update_ui_on_version_or_id,
+        inputs=[version_dropdown, question_input],
+        outputs=[question_box, gt_sql_box, full_table, generated_sql_box, range_md],
     )
 
-    app.load(
-        fn=update_question_info,
-        inputs=[question_input],
-        outputs=[question_box, gt_sql_box, full_table, generated_sql_box],
+    run_button.click(
+        handle_inference,
+        inputs=[version_dropdown, model_dropdown, few_shot_dropdown, question_input],
+        outputs=[generated_sql_box, generated_table, gt_table],
     )
 
-    # 3. Inference Button
-    run_button.click(
-        fn=handle_inference,
-        inputs=[model_dropdown, few_shot_dropdown, question_input],
-        outputs=[
-            generated_sql_box,
-            generated_table,
-            gt_table,
-        ],
+    app.load(
+        update_ui_on_version_or_id,
+        inputs=[version_dropdown, question_input],
+        outputs=[question_box, gt_sql_box, full_table, generated_sql_box, range_md],
     )
 
 app.launch()

dataset_generator.py

@@ -1,19 +1,61 @@
+import json
+
 import pandas as pd
 
-tables = pd.read_pickle("llmsql_tables.pkl")
-questions_splits = pd.read_pickle("llmsql_formatted.pkl")
+
+def load_jsonl(path: str) -> list[dict]:
+    """Load a JSONL file into a list of dicts."""
+    with open(path, encoding="utf-8") as f:
+        return [json.loads(line) for line in f if line.strip()]
+
+
+## LLMSQL 1.0
+tables_llmsql_1 = pd.read_pickle("llmsql_1.0/llmsql_tables.pkl")
+questions_splits_llmsql_1 = pd.read_pickle("llmsql_1.0/llmsql_formatted.pkl")
+
+
+split_info_llmsql_1 = dict()
+questions_llmsql_1 = dict()
+for split_name in list(questions_splits_llmsql_1.keys()):
+    sorted_keys = sorted(
+        list(questions_splits_llmsql_1[split_name].keys()), key=lambda x: int(x)
+    )
+    split_info_llmsql_1[split_name] = {
+        "first": sorted_keys[0],
+        "last": sorted_keys[-1],
+        "count": len(sorted_keys),
+    }
+    for question_id, items in questions_splits_llmsql_1[split_name].items():
+        questions_llmsql_1[int(question_id)] = items
+
+tables_llmsql_1 = pd.read_pickle("llmsql_1.0/llmsql_tables.pkl")
+questions_splits_llmsql_1 = pd.read_pickle("llmsql_1.0/llmsql_formatted.pkl")
 
 
 split_info = dict()
-questions = dict()
-for split_name in list(questions_splits.keys()):
+questions_llmsql_1 = dict()
+for split_name in list(questions_splits_llmsql_1.keys()):
     sorted_keys = sorted(
-        list(questions_splits[split_name].keys()), key=lambda x: int(x)
+        list(questions_splits_llmsql_1[split_name].keys()), key=lambda x: int(x)
     )
     split_info[split_name] = {
         "first": sorted_keys[0],
         "last": sorted_keys[-1],
         "count": len(sorted_keys),
     }
-    for question_id, items in questions_splits[split_name].items():
-        questions[int(question_id)] = items
+    for question_id, items in questions_splits_llmsql_1[split_name].items():
+        questions_llmsql_1[int(question_id)] = items
+
+
+## LLMSQL 2.0
+tables_list_llmsql_2 = load_jsonl("llmsql_2.0/tables.jsonl")
+questions_list_llmsql_2 = load_jsonl("llmsql_2.0/questions.jsonl")
+
+
+questions_llmsql_2 = dict()
+for question in questions_list_llmsql_2:
+    questions_llmsql_2[question["question_id"]] = question
+
+tables_llmsql_2 = dict()
+for table in tables_list_llmsql_2:
+    tables_llmsql_2[table["table_id"]] = table

llmsql_2.0/questions.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dbdaf154133f693142501a807613c3555105d72dd9c9c9cff509063f43efec14
+size 18317052

llmsql_2.0/tables.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b7487109222fba276f843523536d0efd5dd3ab24d6343fd84d19fe94aba3eeb
+size 45341094