src/bird/inference.py

"""
Shared inference utilities for BIRD pipelines.

Provides:
- build_instruction(): the canonical prompt format (must match training)
- BIRDvLLMEngine: vLLM-backed batched generation with optional Best-of-N
- _flash_attn_available(): helper used by trainer scripts

This module is imported by:
  - src.bird.eval_finetuned         (build_instruction)
  - src.bird.build_sft_data         (build_instruction)
  - src.bird.build_pairs            (build_instruction)
  - src.bird.sft_train              (_flash_attn_available)
  - src.bird.dpo_train              (_flash_attn_available)
  - src.bird.eval_best_of_n         (BIRDvLLMEngine)
  - src.bird.build_self_sampling_pairs (BIRDvLLMEngine)
"""
from __future__ import annotations

import re
from typing import Iterable


# ─────────────────────────────────────────────────────────────────────────────
# PROMPT — must match exactly what was used during DPO training
# Single source of truth across the codebase
# ─────────────────────────────────────────────────────────────────────────────

def build_instruction(question: str, schema: str, evidence: str = "") -> str:
    """Canonical BIRD prompt format. Used by training data builders + eval."""
    evidence_block = f"External Knowledge:\n{evidence}\n\n" if evidence.strip() else ""
    return (
        "Convert the following natural language question into a valid SQL query.\n\n"
        f"Database Schema:\n{schema}\n\n"
        f"{evidence_block}"
        f"Question: {question}\n\n"
        "Return only the SQL query with no explanation."
    )


SYSTEM_PROMPT = (
    "You are an expert SQL query generator. "
    "Given a database schema and a natural language question, "
    "generate a single valid SQL query that answers the question. "
    "Output ONLY the SQL query, nothing else."
)


# ─────────────────────────────────────────────────────────────────────────────
# OUTPUT POSTPROCESSING — strip markdown fences, add trailing semicolon
# ─────────────────────────────────────────────────────────────────────────────

_FENCE_RE = re.compile(r"```(?:sql)?\s*\n?(.*?)```", re.DOTALL)


def extract_sql(text: str) -> str:
    """Strip markdown fences, return SQL with trailing semicolon."""
    text  = text.strip()
    match = _FENCE_RE.search(text)
    if match:
        text = match.group(1).strip()
    return text if text.endswith(";") else text + ";"


# ─────────────────────────────────────────────────────────────────────────────
# CHAIN-OF-THOUGHT PROMPT — for the CoT-SFT model (qwen2.5-coder-7b-bird-cot).
# The model REASONS first, then emits the final query in a ```sql block.
# Kept separate from build_instruction/extract_sql so the direct-SQL path is
# untouched. Must match the CoT-SFT training prompt exactly.
# ─────────────────────────────────────────────────────────────────────────────

COT_SYSTEM_PROMPT = (
    "You are an expert SQLite query writer. Reason step by step about the schema "
    "and the question, then output the final query in a ```sql code block."
)


def build_cot_instruction(question: str, schema: str, evidence: str = "") -> str:
    """CoT prompt: ask for step-by-step reasoning, then a final fenced SQL query."""
    evidence_block = f"External Knowledge:\n{evidence}\n\n" if evidence.strip() else ""
    return (
        "Given the database schema and question, work out the correct SQLite "
        "query step by step.\n\n"
        f"Database Schema:\n{schema}\n\n"
        f"{evidence_block}"
        f"Question: {question}\n\n"
        "Think step by step:\n"
        "1. Which tables and columns are relevant?\n"
        "2. What joins, filters, grouping, and ordering are needed?\n"
        "3. Apply any external knowledge and handle edge cases.\n\n"
        "Then give the final answer as:\n"
        "```sql\n<final query>\n```"
    )


_COT_FENCE_RE = re.compile(r"```(?:sql)?\s*\n?(.*?)```", re.DOTALL | re.IGNORECASE)


def extract_final_sql(text: str) -> str:
    """Return the LAST fenced SQL block — the final query after the reasoning.

    (Differs from extract_sql, which takes the FIRST fence: in CoT output the
    model may show intermediate queries while reasoning; we want the last one.)
    """
    text = text.strip()
    matches = _COT_FENCE_RE.findall(text)
    if matches:
        sql = matches[-1].strip()
    else:
        # No code fence — fall back to the last SELECT/WITH statement onward.
        starts = list(re.finditer(r"(?is)\b(WITH|SELECT)\b", text))
        sql = text[starts[-1].start():].strip() if starts else text
    return sql if sql.endswith(";") else sql + ";"


# ─────────────────────────────────────────────────────────────────────────────
# FLASH-ATTENTION HELPER — used by SFT + DPO trainers
# ─────────────────────────────────────────────────────────────────────────────

def _flash_attn_available() -> bool:
    try:
        import flash_attn  # noqa: F401
        return True
    except ImportError:
        return False


# ─────────────────────────────────────────────────────────────────────────────
# vLLM BATCHED INFERENCE ENGINE
# ─────────────────────────────────────────────────────────────────────────────

class BIRDvLLMEngine:
    """
    vLLM-backed batched SQL generation. Supports K-sample generation for
    Best-of-N inference and self-sampling preference pair generation.

    Loads base model + (optional) LoRA adapter once. All subsequent
    generate() calls reuse the loaded engine.

    Usage:
        engine = BIRDvLLMEngine(
            base_model="Qwen/Qwen2.5-Coder-14B-Instruct",
            adapter_path="/scratch/.../bird_sft_adapter_14b/checkpoint-3100",
        )
        # Each item is a tuple (question, schema, evidence)
        prompts = [(q, s, e) for q, s, e in items]
        # Returns: list of (list of K candidate SQL strings)
        candidates_per_question = engine.generate(prompts, k=4, temperature=0.8)
    """

    def __init__(
        self,
        base_model: str,
        adapter_path: str | None = None,
        max_model_len: int = 8192,
        gpu_memory_utilization: float = 0.90,
        dtype: str = "bfloat16",
    ):
        from vllm import LLM
        from vllm.lora.request import LoRARequest
        from transformers import AutoTokenizer

        self.tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        self.llm = LLM(
            model=base_model,
            dtype=dtype,
            max_model_len=max_model_len,
            gpu_memory_utilization=gpu_memory_utilization,
            enable_lora=adapter_path is not None,
            max_lora_rank=64,  # safe upper bound; covers r=32 with margin
            trust_remote_code=True,
        )

        self.lora_request = None
        if adapter_path:
            self.lora_request = LoRARequest(
                lora_name="bird_adapter",
                lora_int_id=1,
                lora_path=adapter_path,
            )
            print(f"vLLM engine loaded with adapter: {adapter_path}")
        else:
            print(f"vLLM engine loaded (base model only, no adapter)")

    def _build_chat_prompt(self, question: str, schema: str, evidence: str) -> str:
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": build_instruction(question, schema, evidence)},
        ]
        return self.tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

    def generate(
        self,
        items: Iterable[tuple[str, str, str]],
        k: int = 1,
        temperature: float = 0.0,
        top_p: float = 1.0,
        max_tokens: int = 512,
    ) -> list[list[str]]:
        """
        Args:
            items: iterable of (question, schema, evidence) tuples
            k: number of candidates per question (1 = greedy, K>1 = sampling)
            temperature: 0.0 for greedy, 0.7-1.0 for sampling
            max_tokens: max generated tokens per candidate

        Returns:
            list of length len(items); each element is a list of K SQL strings
        """
        from vllm import SamplingParams

        items = list(items)
        prompts = [self._build_chat_prompt(q, s, e) for q, s, e in items]

        sampling_params = SamplingParams(
            n=k,
            temperature=temperature if k > 1 else 0.0,
            top_p=top_p,
            max_tokens=max_tokens,
        )

        outputs = self.llm.generate(
            prompts,
            sampling_params=sampling_params,
            lora_request=self.lora_request,
        )

        results: list[list[str]] = []
        for out in outputs:
            cands = [extract_sql(o.text) for o in out.outputs]
            results.append(cands)
        return results