src/text2sql.py

"""
Local text-to-SQL over SQLite databases.

The public product path is:
  - load the merged CoT-SFT model (qwen2.5-coder-7b-bird-cot), or a base
    model plus a SQL LoRA adapter
  - introspect a SQLite schema
  - reason step by step and generate SQL from a natural-language question
  - optionally execute only read-only SQL against the database
"""
from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path
from typing import Any

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

from src.bird.inference import (
    build_instruction,
    build_cot_instruction,
    extract_sql,
    extract_final_sql,
    COT_SYSTEM_PROMPT,
)
from src.shared.schema_loader import get_schema_from_sqlite
from src.shared.sqlite_executor import execute_sqlite_query

# Merged chain-of-thought model — the honest, uncontaminated result. Trained via
# CoT-SFT on BIRD (train split), evaluated on the BIRD dev split: 52.1% greedy /
# 58.5% self-consistency (K=8). Loaded directly, no separate adapter needed.
DEFAULT_MODEL = "jk200201/qwen2.5-coder-7b-bird-cot"
DEFAULT_BASE = "Qwen/Qwen2.5-Coder-7B-Instruct"

# CLI adapter shortcuts (base + LoRA). The hosted demo uses the merged CoT model.
# The spider/bird entries are legacy DPO adapters kept for reproducibility only;
# the spider one has a contaminated dev eval and should not be quoted for scores.
ADAPTERS = {
    "cot":    "jk200201/qwen2.5-coder-7b-bird-cot-lora",
    "spider": "jk200201/qwen2.5-coder-7b-sql-dpo",
    "bird":   "jk200201/qwen2.5-coder-7b-bird-dpo",
    "base":   "",
}


def resolve_adapter(adapter: str | None) -> str | None:
    """Accept a shortcut, HF repo id, local path, empty string, or None.

    None means "no adapter" — load the merged DEFAULT_MODEL directly.
    """
    if adapter is None:
        return None
    if adapter in ADAPTERS:
        return ADAPTERS[adapter] or None
    return adapter or None


def load_model(
    model_id: str = DEFAULT_MODEL,
    adapter: str | None = None,
    use_4bit: bool = True,
):
    """Load a model for inference.

    If `adapter` is given, load DEFAULT_BASE + that LoRA adapter. Otherwise load
    `model_id` directly (a merged/full model, e.g. the CoT-SFT model).
    """
    adapter = resolve_adapter(adapter)

    if use_4bit and not torch.cuda.is_available():
        raise RuntimeError(
            "4-bit inference needs a CUDA GPU. On Hugging Face Spaces, open "
            "Settings -> Hardware and choose a GPU such as Nvidia L4 before "
            "running the demo."
        )

    # When using an adapter the weights come from the base; otherwise from model_id.
    weights = DEFAULT_BASE if adapter else model_id

    print(
        f"Loading {weights}" + (f" + {adapter}" if adapter else " (merged)"),
        file=sys.stderr,
    )
    tokenizer = AutoTokenizer.from_pretrained(weights, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    kwargs: dict[str, Any] = {"trust_remote_code": True, "device_map": "auto"}
    if use_4bit:
        kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
    else:
        kwargs["torch_dtype"] = torch.bfloat16

    model = AutoModelForCausalLM.from_pretrained(weights, **kwargs)
    if adapter:
        model = PeftModel.from_pretrained(model, adapter)
    model.eval()
    return model, tokenizer


def generate_sql(
    model,
    tokenizer,
    schema: str,
    question: str,
    evidence: str = "",
    cot: bool = True,
    max_new_tokens: int = 1024,
) -> tuple[str, str]:
    """Build the training-time prompt, greedily decode, and return (sql, reasoning).

    CoT mode (default): system + step-by-step prompt, big token budget, and the
    final SQL is the LAST fenced block. `reasoning` is the full generation.
    Direct mode: the legacy "return only SQL" prompt; `reasoning` is empty.
    """
    if cot:
        messages = [
            {"role": "system", "content": COT_SYSTEM_PROMPT},
            {"role": "user", "content": build_cot_instruction(question, schema, evidence)},
        ]
    else:
        messages = [{"role": "user", "content": build_instruction(question, schema, evidence)}]

    # return_dict=True + unpacking is version-robust: apply_chat_template returns a
    # BatchEncoding (dict) in newer transformers, and passing that positionally to
    # generate() makes it treat the dict as input_ids and fail on `.shape`.
    enc = tokenizer.apply_chat_template(
        messages,
        return_tensors="pt",
        add_generation_prompt=True,
        return_dict=True,
    )
    enc = {k: v.to(model.device) for k, v in enc.items()}
    input_len = enc["input_ids"].shape[-1]

    with torch.no_grad():
        outputs = model.generate(
            **enc,
            max_new_tokens=max_new_tokens if cot else min(max_new_tokens, 256),
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )
    raw = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True)

    if cot:
        return extract_final_sql(raw), raw.strip()
    return extract_sql(raw), ""


def is_read_only_sql(sql: str) -> bool:
    """
    Conservative guard for product execution.

    We allow SELECT/WITH queries only and reject obvious mutating statements.
    This keeps the CLI/demo from modifying a user's database by accident.
    """
    cleaned = re.sub(r"--.*?$|/\*.*?\*/", "", sql, flags=re.MULTILINE | re.DOTALL).strip()
    if not cleaned:
        return False
    first_token = cleaned.split(None, 1)[0].lower().rstrip(";")
    if first_token not in {"select", "with"}:
        return False
    blocked = re.search(
        r"\b(insert|update|delete|drop|alter|create|replace|attach|detach|vacuum|pragma)\b",
        cleaned,
        flags=re.IGNORECASE,
    )
    return blocked is None


def predict(
    db_path: str,
    question: str,
    model,
    tokenizer,
    evidence: str = "",
    execute: bool = True,
    cot: bool = True,
    max_new_tokens: int = 1024,
) -> dict:
    """Introspect schema, reason + generate SQL, and optionally run it read-only."""
    schema = get_schema_from_sqlite(db_path)
    sql, reasoning = generate_sql(
        model, tokenizer, schema, question, evidence, cot=cot, max_new_tokens=max_new_tokens
    )

    result = {
        "sql": sql,
        "reasoning": reasoning,
        "columns": [],
        "rows": [],
        "row_count": 0,
        "error": None,
        "schema": schema,
    }

    if execute:
        if not is_read_only_sql(sql):
            result["error"] = "Refusing to execute non-read-only SQL. Use --no-exec to inspect it."
            return result
        exec_out = execute_sqlite_query(sql, db_path)
        result.update(
            columns=exec_out["columns"],
            rows=exec_out["rows"],
            row_count=exec_out["row_count"],
            error=exec_out["error"],
        )
    return result


def _print_result(result: dict) -> None:
    if result.get("reasoning"):
        print("\nReasoning")
        print(result["reasoning"])

    print("\nSQL")
    print(result["sql"])

    if result["error"]:
        print("\nError")
        print(result["error"])
        return

    print(f"\nResults ({result['row_count']} rows)")
    try:
        from tabulate import tabulate

        print(tabulate(result["rows"][:50], headers=result["columns"], tablefmt="github"))
    except ImportError:
        print(result["columns"])
        for row in result["rows"][:50]:
            print(row)
    if result["row_count"] > 50:
        print(f"... ({result['row_count'] - 50} more rows)")


def main() -> None:
    parser = argparse.ArgumentParser(description="Local text-to-SQL over a SQLite database")
    parser.add_argument("--db", required=True, help="Path to a .sqlite database file")
    parser.add_argument("--q", "--question", dest="question", required=True)
    parser.add_argument(
        "--model",
        default=DEFAULT_MODEL,
        help="Merged model repo id or local path (default: the CoT-SFT model)",
    )
    parser.add_argument(
        "--adapter",
        default=None,
        help="Optional LoRA adapter on top of the base: repo id, path, or shortcut "
             "(cot, spider, bird, base). Omit to use the merged --model directly.",
    )
    parser.add_argument("--evidence", default="", help="Optional BIRD-style domain hint")
    parser.add_argument("--bf16", action="store_true", help="Load in bf16 instead of 4-bit")
    parser.add_argument("--direct", action="store_true", help="Direct-SQL prompt (no reasoning)")
    parser.add_argument("--no-exec", action="store_true", help="Generate SQL without running it")
    parser.add_argument("--max-new-tokens", type=int, default=1024)
    args = parser.parse_args()

    db_path = Path(args.db)
    if not db_path.exists():
        sys.exit(f"Database not found: {db_path}")

    model, tokenizer = load_model(
        model_id=args.model,
        adapter=args.adapter,
        use_4bit=not args.bf16,
    )
    result = predict(
        str(db_path),
        args.question,
        model,
        tokenizer,
        evidence=args.evidence,
        execute=not args.no_exec,
        cot=not args.direct,
        max_new_tokens=args.max_new_tokens,
    )
    _print_result(result)


if __name__ == "__main__":
    main()