"""Model-backed evaluation policy for SQLEnv. This is the missing piece the success gate ("beat 28-32% on the N=50 Spider subset") depends on: a `Policy` that drives the TRAINED LLM through the environment, so `evaluate(env, ModelPolicy(...), n_episodes=...)` produces a real `success_rate`. The same agentic-rollout primitive is what the Gradio Space needs to serve a question, so build/validate it once here. IMPORTANT — this reconstructs the multi-turn tool-use transcript INDEPENDENTLY of TRL. During training, TRL's `environment_factory` owns the chat rendering and tool dispatch; that machinery is internal and not reusable at eval time. So this policy maintains its own message history and parses the model's emitted tool call itself. That introduces a small train/serve-skew risk in prompt formatting — **validate it** by running the smoke eval and confirming the parse rate is ~95-100% (matching training) before trusting any accuracy number. Heavy deps (torch/transformers) are imported lazily so this module lint-imports without the training extras. `ModelPolicy` satisfies the `Policy` protocol in ``evaluation.policies`` (``select_action(observation) -> SQLAction``). """ from __future__ import annotations import logging from typing import Any try: from sql_env.models import SQLAction, SQLObservation except ImportError: # pragma: no cover - flat-layout / Docker fallback from models import SQLAction, SQLObservation # type: ignore[no-redef] try: from sql_env.server.tooling import ( get_system_prompt, get_tool_definitions, parse_action, ) except ImportError: # pragma: no cover - flat-layout / Docker fallback from server.tooling import ( # type: ignore[no-redef] get_system_prompt, get_tool_definitions, parse_action, ) logger = logging.getLogger(__name__) class ModelPolicy: """Drive a trained causal-LM through SQLEnv as an agentic tool-use policy. Parameters ---------- model, tokenizer A loaded causal LM + tokenizer (base+adapter already merged/attached). enable_thinking Must match how the model was trained (default False -> /no_think). temperature, top_p, top_k Inference sampling. Defaults are Qwen3 non-thinking inference settings (T=0.7, top_p=0.8, top_k=20) — NOT the RL training temperature (1.0). max_new_tokens Per-turn generation cap. """ def __init__( self, model: Any, tokenizer: Any, *, enable_thinking: bool = False, temperature: float = 0.7, top_p: float = 0.8, top_k: int = 20, max_new_tokens: int = 512, ) -> None: self._model = model self._tokenizer = tokenizer self._enable_thinking = enable_thinking self._temperature = temperature self._top_p = top_p self._top_k = top_k self._max_new_tokens = max_new_tokens self._system_prompt = get_system_prompt(enable_thinking=enable_thinking) self._tools = get_tool_definitions() # Per-episode conversation state, (re)initialised when step_count == 0. self._messages: list[dict[str, str]] = [] self._started = False def select_action(self, observation: SQLObservation) -> SQLAction: """Choose one tool call for the current observation.""" # A fresh episode: evaluate() calls env.reset() (step_count == 0) before # the first select_action. Detect that and reset our transcript. if observation.step_count == 0 or not self._started: self._messages = [ {"role": "system", "content": self._system_prompt}, { "role": "user", "content": ( f"{observation.question}\n\n" f"Available tables:\n{observation.schema_info}" ), }, ] self._started = True else: # The observation carries the result/error of OUR previous tool call. tool_response = observation.result or ( f"Error: {observation.error}" if observation.error else "No output." ) self._messages.append({"role": "tool", "content": tool_response}) completion = self._generate(self._messages) self._messages.append({"role": "assistant", "content": completion}) action = self._parse_action(completion) if action is None: # Unparseable output near the budget: submit a best-effort answer so # the episode terminates instead of burning steps on malformed calls. if observation.budget_remaining <= 1: return SQLAction(action_type="ANSWER", argument="[]") # Otherwise emit a harmless exploration step; the parse failure is # itself the signal we watch (parse rate). return SQLAction(action_type="QUERY", argument="SELECT 1") return action def _generate(self, messages: list[dict[str, str]]) -> str: import torch # noqa: PLC0415 inputs = self._tokenizer.apply_chat_template( messages, tools=self._tools, add_generation_prompt=True, return_tensors="pt", tokenize=True, # transformers v5 flipped the default to return a BatchEncoding # dict; this loop (and the trained format) expects the bare # input_ids tensor — pin the pre-v5 behavior explicitly. return_dict=False, ) inputs = inputs.to(self._model.device) # Pass an explicit all-ones attention mask. We generate one unpadded # sequence at a time, so this is behaviorally a no-op today — but because # Qwen3's pad_token == eos_token, transformers cannot auto-infer the mask # and warns ("attention mask is not set ... you may observe unexpected # behavior"). Passing it silences the warning AND prevents a real # correctness bug if we ever batch eval generation (padded sequences need # a mask, or the model attends to pad tokens). attention_mask = torch.ones_like(inputs) with torch.no_grad(): out = self._model.generate( inputs, attention_mask=attention_mask, max_new_tokens=self._max_new_tokens, do_sample=self._temperature > 0, temperature=self._temperature, top_p=self._top_p, top_k=self._top_k, pad_token_id=self._tokenizer.eos_token_id, ) generated = out[0][inputs.shape[-1] :] return self._tokenizer.decode(generated, skip_special_tokens=False) @staticmethod def _parse_action(completion: str) -> SQLAction | None: """Extract the first JSON and map it to an SQLAction. Thin back-compat alias delegating to ``sql_env.server.tooling.parse_action``. """ return parse_action(completion)