ritvik360/nl2sql-bench

1"""2inference.py  —  NL2SQL-Bench Baseline Inference Script3========================================================4 5MANDATORY COMPLIANCE6--------------------7- Named `inference.py`, placed in project root.8- Uses OpenAI client for all LLM calls.9- Reads: API_BASE_URL, MODEL_NAME, HF_TOKEN from environment.10- Emits [START] / [STEP] / [END] lines to stdout in the exact format below.11- Runs all 3 tasks; total runtime < 20 min on 2 vCPU / 8 GB.12 13STDOUT FORMAT (exact — any deviation breaks scoring)14----------------------------------------------------15[START] task=<task_name> env=nl2sql-bench model=<model_name>16[STEP]  step=<n> action=<sql_one_line> reward=<0.00> done=<true|false> error=<msg|null>17[END]   success=<true|false> steps=<n> score=<0.000> rewards=<r1,r2,...>18"""19 20from __future__ import annotations21 22import asyncio23import os24import sys25import textwrap26from typing import List, Optional27 28from openai import OpenAI29 30# # ── Configuration ──────────────────────────────────────────────────────────31# API_BASE_URL   = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")32# MODEL_NAME     = os.getenv("MODEL_NAME",   "Qwen/Qwen2.5-7B-Instruct")33# API_KEY        = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY", "")34# IMAGE_NAME     = os.getenv("LOCAL_IMAGE_NAME", "nl2sql-bench:latest")35# SPACE_URL      = os.getenv("SPACE_URL", "http://localhost:8000")36 37# BENCHMARK      = "nl2sql-bench"38# MAX_STEPS      = 539# TEMPERATURE    = 0.2      # Low temp for SQL generation40# MAX_TOKENS     = 51241# SUCCESS_THRESHOLD = 0.7   # score >= 0.7 → success42 43# TASKS = ["simple-filter", "join-aggregation", "analytics-window"]44 45# ── Configuration ──────────────────────────────────────────────────────────46API_BASE_URL      = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")47# Points to your newly uploaded fine-tuned weights!48MODEL_NAME        = os.getenv("MODEL_NAME", "ritvik360/qwen-7b-nl2sql-merged_1") 49# CRITICAL FIX: Looks for 'API_KEY' first to satisfy the evaluator's LiteLLM proxy50API_KEY           = os.getenv("API_KEY")  or os.getenv("HF_TOKEN", "") or os.getenv("OPENAI_API_KEY")51IMAGE_NAME        = os.getenv("LOCAL_IMAGE_NAME", "nl2sql-bench:latest")52# CRITICAL FIX: Point the default directly to your live HF Space!53SPACE_URL         = os.getenv("SPACE_URL", "https://ritvik360-nl2sql-bench.hf.space")54 55BENCHMARK         = "nl2sql-bench"56MAX_STEPS         = 557TEMPERATURE       = 0.2      # Low temp for SQL generation58MAX_TOKENS        = 51259SUCCESS_THRESHOLD = 0.7      # score >= 0.7 → success60 61TASKS = ["simple-filter", "join-aggregation", "analytics-window"]62 63# ── System prompt ──────────────────────────────────────────────────────────64SYSTEM_PROMPT = textwrap.dedent("""65You are an expert SQL analyst working with a SQLite e-commerce database.66 67DATABASE SCHEMA68---------------69categories(id, name)70products(id, name, category_id, price, stock_quantity)71customers(id, name, email, country, tier∈{bronze|silver|gold}, created_at)72orders(id, customer_id, status∈{pending|processing|shipped|delivered|cancelled},73       created_at, total_amount)74order_items(id, order_id, product_id, quantity, unit_price)75reviews(id, product_id, customer_id, rating∈1-5, created_at)76 77RULES78-----791. Write a single SELECT query — no INSERT/UPDATE/DELETE.802. Output ONLY the SQL query, nothing else. No markdown, no explanation.813. Use SQLite syntax: strftime('%Y-%m', date_col) for month, ROUND(x, 2) for decimals.824. Window functions (RANK, DENSE_RANK, ROW_NUMBER, running SUM) are supported.835. CTEs (WITH ... AS (...)) are supported.846. If you receive an error, fix it carefully in your next attempt.857. If you receive partial results, refine your query to match the expected output.86""").strip()87 88 89# ── Stdout logging (mandatory format) ─────────────────────────────────────90 91def log_start(task: str, model: str) -> None:92    print(f"[START] task={task} env={BENCHMARK} model={model}", flush=True)93 94 95def log_step(96    step: int, action: str, reward: float, done: bool, error: Optional[str]97) -> None:98    # Collapse multi-line SQL to single line for log compliance99    action_single = " ".join(action.split())100    error_val = error.replace("\n", " ") if error else "null"101    print(102        f"[STEP] step={step} action={action_single!r} "103        f"reward={reward:.2f} done={str(done).lower()} error={error_val}",104        flush=True,105    )106 107 108def log_end(109    success: bool, steps: int, score: float, rewards: List[float]110) -> None:111    rewards_str = ",".join(f"{r:.2f}" for r in rewards)112    print(113        f"[END] success={str(success).lower()} steps={steps} "114        f"score={score:.3f} rewards={rewards_str}",115        flush=True,116    )117 118 119# ── LLM interaction ────────────────────────────────────────────────────────120 121def build_user_prompt(122    question: str,123    schema_context: str,124    step: int,125    last_query: str,126    last_error: Optional[str],127    last_result: list,128    result_columns: list,129) -> str:130    parts = [f"QUESTION: {question}", ""]131 132    if step > 1:133        parts.append(f"Your previous SQL (step {step - 1}):")134        parts.append(f"  {' '.join(last_query.split())}")135        parts.append("")136        if last_error:137            parts.append(f"ERROR: {last_error}")138        elif last_result:139            preview = str(last_result[:3]).replace("\n", " ")140            parts.append(f"RESULT PREVIEW (first 3 rows): {preview}")141            parts.append(f"COLUMNS: {result_columns}")142        parts.append("")143        parts.append("Please correct or refine your query.")144    else:145        parts.append("Write a SQL query to answer the question.")146 147    return "\n".join(parts)148 149 150def call_llm(client: OpenAI, user_prompt: str) -> str:151    try:152        resp = client.chat.completions.create(153            model=MODEL_NAME,154            messages=[155                {"role": "system", "content": SYSTEM_PROMPT},156                {"role": "user",   "content": user_prompt},157            ],158            temperature=TEMPERATURE,159            max_tokens=MAX_TOKENS,160            stream=False,161        )162        text = (resp.choices[0].message.content or "").strip()163        # Strip markdown code fences if model wraps in ```sql ... ```164        if text.startswith("```"):165            lines = text.split("\n")166            text = "\n".join(167                l for l in lines168                if not l.strip().startswith("```")169            ).strip()170        return text if text else "SELECT 1"171    except Exception as exc:172        print(f"[DEBUG] LLM call failed: {exc}", file=sys.stderr, flush=True)173        return "SELECT 1"174 175 176# ── Single-task episode ────────────────────────────────────────────────────177 178async def run_task(client: OpenAI, env, task_name: str) -> dict:179    """Run one full episode for the given task. Returns result dict."""180    rewards: List[float] = []181    steps_taken = 0182    score = 0.0183    success = False184 185    log_start(task_name, MODEL_NAME)186 187    try:188        # Reset — pass task_name via action payload or query param189        # OpenEnv reset() may not accept task args via HTTP; we rely on190        # NL2SQL_DEFAULT_TASK env-var being set before calling, OR we191        # pass it as a reset parameter if the server supports it.192        result = await env.reset() # changed193        obs = result.observation194 195        for step in range(1, MAX_STEPS + 1):196            if result.done:197                break198 199            user_prompt = build_user_prompt(200                question=obs.question,201                schema_context=obs.schema_context,202                step=step,203                last_query=obs.last_query,204                last_error=obs.last_error,205                last_result=obs.last_result,206                result_columns=obs.result_columns,207            )208 209            sql = call_llm(client, user_prompt)210 211            from models import NL2SQLAction  # local to avoid circular at module level212            action = NL2SQLAction(query=sql)213            result = await env.step(action)214            obs = result.observation215 216            reward = obs.reward or 0.0217            done   = obs.done218            error  = obs.last_error219 220            rewards.append(reward)221            steps_taken = step222 223            log_step(step=step, action=sql, reward=reward, done=done, error=error)224 225            if done:226                break227 228        # Compute final score229        # CRITICAL: Evaluator requires score strictly in (0, 1) — not 0.0, not 1.0.230        # A perfect solve gives 1.0 → clamp to 0.999. All-fail gives 0.0 → clamp to 0.001.231        raw_score = sum(rewards) / max(len(rewards), 1)232        score     = round(min(max(raw_score, 0.001), 0.999), 4)233        success   = raw_score >= SUCCESS_THRESHOLD234 235    except Exception as exc:236        print(f"[DEBUG] Episode error for {task_name}: {exc}", file=sys.stderr, flush=True)237    finally:238        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)239 240    return {"task": task_name, "success": success, "score": score, "rewards": rewards}241 242 243# ── Main ───────────────────────────────────────────────────────────────────244 245async def main() -> None:246    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)247 248    # Import here to avoid import errors if openenv not installed during lint249    from client import NL2SQLEnv250 251    all_results = []252 253    for task_name in TASKS:254        # Set the default task for the server session via env-var approach.255        # For the hosted Space, we rely on the task cycling implemented in256        # the task registry's round-robin iterator.257        os.environ["NL2SQL_DEFAULT_TASK"] = task_name258 259        try:260            async with NL2SQLEnv(base_url=SPACE_URL) as env:261                result = await run_task(client, env, task_name)262                all_results.append(result)263        except Exception as exc:264            print(265                f"[DEBUG] Failed to connect for task {task_name}: {exc}",266                file=sys.stderr,267                flush=True,268            )269            # Emit a zero-score END to keep log format valid270            log_end(success=False, steps=0, score=0.0, rewards=[])271            all_results.append({"task": task_name, "success": False, "score": 0.0})272 273    # Summary to stderr (not scored, for human readability)274    print("\n=== Baseline Summary ===", file=sys.stderr)275    for r in all_results:276        print(277            f"  {r['task']:20s}  score={r['score']:.3f}  "278            f"success={r['success']}",279            file=sys.stderr,280        )281    avg = sum(r["score"] for r in all_results) / max(len(all_results), 1)282    print(f"  {'AVERAGE':20s}  score={avg:.3f}", file=sys.stderr)283 284 285if __name__ == "__main__":286    asyncio.run(main())