scripts/eval_models.py

#!/usr/bin/env python3
"""Evaluate LLMs on the Squirrel-mini toy benchmark via HF Inference Providers.

Mirrors the paper's evaluation prompts (Appendix J.3) and metrics (EM, GM-proxy,
MB-proxy; see scripts/metrics.py for the GM/MB substitutions and why they differ
from the paper's Apache-Calcite-based Graph Match).
"""
from __future__ import annotations

import argparse
import json
import re
import sys
import time
from pathlib import Path

from huggingface_hub import InferenceClient

sys.path.insert(0, str(Path(__file__).resolve().parent))
from metrics import exact_match, graph_match_proxy, modify_better_proxy  # noqa: E402

SYNTAX_PROMPT = """You are an SQL assistant.

## Task
Based on the error messages and table schema, your task is to fix the issue in the SQL and write the correct SQL.
Remember that you can not change any existing comments and SQL code without errors.

## Input Data
The issue SQL: {issue_sql}
Related tables schema: {ddl}
Error Messages: {error_message}

## Output (JSON):
{{"predict_sql": "The fixed SQL."}}
Respond with ONLY the JSON object, no markdown fences, no explanation."""

SEMANTIC_PROMPT = """You are an SQL assistant.

## Task
Based on the user query and input table schema, please fix the bugs in the Issue SQL and write the corresponding correct SQL code.
Remember that you can not change any existing comments and SQL code without errors.

## Input Data
User Query: {user_query}
Related tables schema: {ddl}
Issue SQL: {issue_sql}

## Output (JSON):
{{"predict_sql": "The fixed SQL."}}
Respond with ONLY the JSON object, no markdown fences, no explanation."""


def extract_predict_sql(raw: str) -> str:
    raw = raw.strip()
    raw = re.sub(r"^```(json)?", "", raw).strip()
    raw = re.sub(r"```$", "", raw).strip()
    try:
        obj = json.loads(raw)
        if isinstance(obj, dict) and "predict_sql" in obj:
            return obj["predict_sql"]
    except Exception:
        pass
    m = re.search(r'"predict_sql"\s*:\s*"(.*)"\s*\}', raw, re.S)
    if m:
        return m.group(1).encode().decode("unicode_escape")
    return raw


def call_model(client: InferenceClient, model: str, prompt: str, retries: int = 3) -> str:
    last_err = None
    for attempt in range(retries):
        try:
            resp = client.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": prompt}],
                max_tokens=4000,
                temperature=0,
            )
            return resp.choices[0].message.content or ""
        except Exception as e:
            last_err = e
            time.sleep(2 * (attempt + 1))
    raise RuntimeError(f"model call failed after {retries} tries: {last_err}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tasks", default="data/squirrel_mini_tasks.json")
    ap.add_argument("--out", default="outputs/eval_results.jsonl")
    ap.add_argument(
        "--models",
        nargs="+",
        default=[
            "Qwen/Qwen2.5-Coder-7B-Instruct",
            "Qwen/Qwen2.5-Coder-32B-Instruct",
            "deepseek-ai/DeepSeek-V3-0324",
            "Qwen/Qwen3-235B-A22B-Instruct-2507",
        ],
    )
    ap.add_argument("--limit", type=int, default=None)
    args = ap.parse_args()

    tasks = json.loads(Path(args.tasks).read_text())
    if args.limit:
        tasks = tasks[: args.limit]

    client = InferenceClient()
    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    results = []
    with out_path.open("w") as f:
        for model in args.models:
            for task in tasks:
                if task["type"] == "syntax":
                    prompt = SYNTAX_PROMPT.format(
                        issue_sql=task["issue_sql"],
                        ddl=task["ddl"],
                        error_message=task["error_message"],
                    )
                else:
                    prompt = SEMANTIC_PROMPT.format(
                        user_query=task["user_query"],
                        ddl=task["ddl"],
                        issue_sql=task["issue_sql"],
                    )
                try:
                    raw = call_model(client, model, prompt)
                    pred_sql = extract_predict_sql(raw)
                    em = exact_match(pred_sql, task["reference_sql"])
                    gm = graph_match_proxy(pred_sql, task["reference_sql"])
                    mb = modify_better_proxy(pred_sql, task["reference_sql"], task["issue_sql"])
                    error = None
                except Exception as e:
                    pred_sql, em, gm, mb = "", False, False, None
                    error = str(e)

                rec = {
                    "model": model,
                    "task_id": task["task_id"],
                    "type": task["type"],
                    "em": em,
                    "gm_proxy": gm,
                    "mb_proxy": mb,
                    "pred_sql": pred_sql,
                    "error": error,
                }
                results.append(rec)
                f.write(json.dumps(rec) + "\n")
                f.flush()
                status = "ERR" if error else ("EM" if em else ("GM" if gm else "miss"))
                print(f"{model} | {task['task_id']} | {status}")

    print(f"\nWrote {len(results)} records to {out_path}")


if __name__ == "__main__":
    main()