Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| """LLM-as-a-Verifier Continuous Logit-Expectation Verification Engine for Relational SQL. | |
| Implements mathematical expectation over logprob distributions of scoring tokens, | |
| decomposed multi-criteria SQL verification, and a Probabilistic Pivot Tournament | |
| for efficient O(K * N) test-time candidate re-ranking. | |
| Reference: | |
| 'LLM-as-a-Verifier: A General-Purpose Verification Framework' (July 2026) | |
| arXiv:2607.05391 | |
| """ | |
| import logging | |
| import math | |
| from collections.abc import Callable, Sequence | |
| from dataclasses import dataclass, field | |
| from pathlib import Path | |
| from typing import Any | |
| from omegaconf import DictConfig, OmegaConf | |
| from ask_the_database.rewards.syntax_reward import SQLSyntaxValidator | |
| logger = logging.getLogger(__name__) | |
| DEFAULT_SCORING_TOKENS: dict[str, float] = { | |
| "1": 0.00, | |
| "2": 0.25, | |
| "3": 0.50, | |
| "4": 0.75, | |
| "5": 1.00, | |
| } | |
| class VerifierCriterion: | |
| """A decomposed evaluation criterion for relational SQL queries.""" | |
| name: str | |
| description: str | |
| weight: float = 1.0 | |
| class CandidateScore: | |
| """Detailed score telemetry for a candidate SQL query.""" | |
| candidate_id: str | |
| sql_query: str | |
| overall_score: float | |
| criteria_scores: dict[str, float] = field(default_factory=dict) | |
| scoring_logprobs: dict[str, float] = field(default_factory=dict) | |
| execution_valid: bool = True | |
| metadata: dict[str, Any] = field(default_factory=dict) | |
| def compute_continuous_expectation( | |
| token_logits_or_logprobs: dict[str, float], | |
| scoring_tokens: dict[str, float] | None = None, | |
| temperature: float = 1.0, | |
| is_logprobs: bool = False, | |
| ) -> tuple[float, dict[str, float]]: | |
| """Calculate the continuous mathematical expectation E[S] over scoring tokens. | |
| Formulation: | |
| E[S] = sum_{v in V_S} P(v | x, c) * val(v) | |
| Args: | |
| token_logits_or_logprobs: Mapping from token string to raw logit or logprob. | |
| scoring_tokens: Mapping from token string to normalized value in [0, 1]. | |
| temperature: Softmax temperature parameter. | |
| is_logprobs: True if input values are log probabilities; False if raw logits. | |
| Returns: | |
| tuple of (expected_score, normalized_token_probabilities). | |
| """ | |
| token_map = scoring_tokens or DEFAULT_SCORING_TOKENS | |
| if not token_logits_or_logprobs or not token_map: | |
| return 0.0, {} | |
| # Filter to tokens present in the scoring vocabulary | |
| relevant_tokens: dict[str, float] = {} | |
| for tok, val in token_logits_or_logprobs.items(): | |
| clean_tok = tok.strip() | |
| if clean_tok in token_map: | |
| relevant_tokens[clean_tok] = val | |
| if not relevant_tokens: | |
| return 0.0, {} | |
| temp = max(1e-5, float(temperature)) | |
| # Compute unnormalized probabilities | |
| if is_logprobs: | |
| # Values are log(P). To renormalize over the scoring subset: | |
| max_logp = max(relevant_tokens.values()) | |
| unnorm = {tok: math.exp((lp - max_logp) / temp) for tok, lp in relevant_tokens.items()} | |
| else: | |
| # Values are raw logits | |
| max_logit = max(relevant_tokens.values()) | |
| unnorm = {tok: math.exp((val - max_logit) / temp) for tok, val in relevant_tokens.items()} | |
| total_mass = sum(unnorm.values()) | |
| if total_mass <= 0.0 or math.isnan(total_mass): | |
| return 0.0, {} | |
| probabilities = {tok: mass / total_mass for tok, mass in unnorm.items()} | |
| # Compute expectation: sum(P(v) * val(v)) | |
| expected_score = sum(probabilities[tok] * token_map[tok] for tok in probabilities) | |
| expected_score = max(0.0, min(1.0, expected_score)) | |
| return round(expected_score, 4), probabilities | |
| class SQLVerifier: | |
| """Continuous Logit-Expectation Verification Engine for Text-to-SQL.""" | |
| def __init__( | |
| self, | |
| config: DictConfig | dict[str, Any] | str | Path | None = None, | |
| client: Callable[[str], dict[str, float]] | None = None, | |
| ) -> None: | |
| """Initialize SQLVerifier with criteria and optional scoring client. | |
| Args: | |
| config: Path to rules YAML, DictConfig, or dict. | |
| client: Optional callable (prompt -> token_logits_or_logprobs). | |
| """ | |
| self.config = self._load_config(config) | |
| self.client = client | |
| self.syntax_validator = SQLSyntaxValidator() | |
| # Parse verifier settings | |
| v_cfg = self.config.get("verifier", {}) | |
| self.model_name: str = v_cfg.get("model_name", "Qwen/Qwen2.5-Coder-7B-Instruct") | |
| self.temperature: float = float(v_cfg.get("temperature", 0.0)) | |
| self.pivot_count: int = int(v_cfg.get("pivot_count", 2)) | |
| self.scoring_tokens: dict[str, float] = dict( | |
| v_cfg.get("scoring_tokens", DEFAULT_SCORING_TOKENS) | |
| ) | |
| # Parse decomposed criteria | |
| self.criteria: dict[str, VerifierCriterion] = {} | |
| criteria_cfg = self.config.get("criteria", {}) | |
| total_weight = 0.0 | |
| for name, crit in criteria_cfg.items(): | |
| desc = crit.get("description", "") | |
| wt = float(crit.get("weight", 1.0)) | |
| self.criteria[name] = VerifierCriterion(name=name, description=desc, weight=wt) | |
| total_weight += wt | |
| # Normalize weights if total > 0 | |
| if total_weight > 0.0: | |
| for crit in self.criteria.values(): | |
| crit.weight = crit.weight / total_weight | |
| # Prompt template | |
| self.prompt_template: str = self.config.get( | |
| "prompt_template", | |
| ( | |
| "You are an expert SQL Verification Agent.\n" | |
| "Evaluate whether the candidate SQL query accurately and safely answers the user's question.\n" | |
| "[Database Schema DDL]:\n{schema_ddl}\n\n" | |
| "[User Question]:\n{question}\n\n" | |
| "[Candidate SQL Query]:\n```sql\n{sql_query}\n```\n\n" | |
| "Evaluation Criterion: {criterion_name}\n" | |
| "{criterion_description}\n\n" | |
| "Score from 1 to 5. Respond with only a single digit score:" | |
| ), | |
| ) | |
| def _load_config(self, config: Any) -> dict[str, Any]: | |
| """Load configuration from file or dict.""" | |
| if config is None: | |
| default_path = Path("configs/verification/sql_verifier_rules.yaml") | |
| if default_path.is_file(): | |
| container = OmegaConf.to_container(OmegaConf.load(default_path), resolve=True) | |
| if isinstance(container, dict): | |
| return {str(k): v for k, v in container.items()} | |
| return { | |
| "verifier": {"model_name": "Qwen/Qwen2.5-Coder-7B-Instruct", "pivot_count": 2}, | |
| "criteria": { | |
| "intent_and_filters": {"weight": 0.35, "description": "Filter fidelity"}, | |
| "schema_grounding_and_joins": { | |
| "weight": 0.30, | |
| "description": "Schema presence", | |
| }, | |
| "aggregation_granularity": {"weight": 0.20, "description": "Aggregation level"}, | |
| "result_sanity": {"weight": 0.15, "description": "Result structure"}, | |
| }, | |
| } | |
| if isinstance(config, (str, Path)): | |
| path = Path(config) | |
| if path.is_file(): | |
| container = OmegaConf.to_container(OmegaConf.load(path), resolve=True) | |
| if isinstance(container, dict): | |
| return {str(k): v for k, v in container.items()} | |
| raise ValueError(f"Expected dict configuration from {path}") | |
| raise FileNotFoundError(f"Configuration file not found: {path}") | |
| if isinstance(config, DictConfig): | |
| container = OmegaConf.to_container(config, resolve=True) | |
| if isinstance(container, dict): | |
| return {str(k): v for k, v in container.items()} | |
| raise ValueError("Expected dict configuration from DictConfig") | |
| if isinstance(config, dict): | |
| return {str(k): v for k, v in config.items()} | |
| raise TypeError(f"Unsupported config type: {type(config)}") | |
| def build_verification_prompt( | |
| self, | |
| question: str, | |
| schema_ddl: str, | |
| sql_query: str, | |
| criterion: VerifierCriterion, | |
| ) -> str: | |
| """Construct the evaluation prompt for a specific decomposed criterion.""" | |
| return self.prompt_template.format( | |
| schema_ddl=schema_ddl.strip(), | |
| question=question.strip(), | |
| sql_query=sql_query.strip(), | |
| criterion_name=criterion.name, | |
| criterion_description=criterion.description, | |
| ) | |
| def score_criterion( | |
| self, | |
| sql_query: str, | |
| question: str, | |
| schema_ddl: str, | |
| criterion: VerifierCriterion, | |
| ) -> tuple[float, dict[str, float]]: | |
| """Score a candidate query along a single criterion using logit expectation.""" | |
| if self.client is not None: | |
| prompt = self.build_verification_prompt(question, schema_ddl, sql_query, criterion) | |
| try: | |
| logits_or_logprobs = self.client(prompt) | |
| return compute_continuous_expectation( | |
| logits_or_logprobs, | |
| scoring_tokens=self.scoring_tokens, | |
| temperature=max(1.0, self.temperature), | |
| ) | |
| except Exception as e: | |
| logger.warning("Verifier client failed for criterion %s: %s", criterion.name, e) | |
| # Fallback to deterministic heuristic when client is None or fails | |
| score = self.fallback_heuristic_score(sql_query, question, schema_ddl, criterion.name) | |
| synthetic_probs = dict.fromkeys(self.scoring_tokens, 0.0) | |
| closest_token = min( | |
| self.scoring_tokens.keys(), | |
| key=lambda t: abs(self.scoring_tokens[t] - score), | |
| ) | |
| synthetic_probs[closest_token] = 1.0 | |
| return score, synthetic_probs | |
| def score_query( | |
| self, | |
| sql_query: str, | |
| question: str, | |
| schema_ddl: str, | |
| candidate_id: str = "c0", | |
| execution_valid: bool = True, | |
| metadata: dict[str, Any] | None = None, | |
| ) -> CandidateScore: | |
| """Evaluate a candidate query across all decomposed criteria and compute expectation.""" | |
| criteria_scores: dict[str, float] = {} | |
| all_logprobs: dict[str, float] = {} | |
| total_score = 0.0 | |
| for name, crit in self.criteria.items(): | |
| s_crit, lprobs = self.score_criterion(sql_query, question, schema_ddl, crit) | |
| criteria_scores[name] = s_crit | |
| all_logprobs.update({f"{name}_{k}": v for k, v in lprobs.items()}) | |
| total_score += s_crit * crit.weight | |
| # Penalize if known execution failed | |
| if not execution_valid: | |
| total_score = total_score * 0.2 | |
| overall = round(max(0.0, min(1.0, total_score)), 4) | |
| return CandidateScore( | |
| candidate_id=candidate_id, | |
| sql_query=sql_query, | |
| overall_score=overall, | |
| criteria_scores=criteria_scores, | |
| scoring_logprobs=all_logprobs, | |
| execution_valid=execution_valid, | |
| metadata=metadata or {}, | |
| ) | |
| def fallback_heuristic_score( | |
| self, | |
| sql_query: str, | |
| question: str, | |
| schema_ddl: str, | |
| criterion_name: str | None = None, | |
| ) -> float: | |
| """Deterministic heuristic fallback when external verifier model is unavailable. | |
| Evaluates AST syntax validity, schema grounding, table joins, and keyword matching. | |
| """ | |
| # 1. AST Validation | |
| syntax_res = self.syntax_validator.validate(sql_query) | |
| if not syntax_res.is_valid: | |
| return 0.10 | |
| parsed_query = sql_query.lower() | |
| q_lower = question.lower() | |
| ddl_lower = schema_ddl.lower() | |
| score = 0.50 # Base score for valid syntax | |
| # 2. Schema Grounding Check | |
| tables = list(syntax_res.referenced_tables) | |
| grounded_tables = [t for t in tables if t in ddl_lower] | |
| if tables and len(grounded_tables) == len(tables): | |
| score += 0.20 | |
| elif tables and not grounded_tables: | |
| score -= 0.30 | |
| # 3. Intent & Criteria Specific Adjustments | |
| if criterion_name == "intent_and_filters" or criterion_name is None: | |
| if "where" in parsed_query: | |
| score += 0.15 | |
| if ( | |
| any(k in q_lower for k in ["how many", "count", "number of"]) | |
| and "count" in parsed_query | |
| ): | |
| score += 0.15 | |
| if ( | |
| (criterion_name == "schema_grounding_and_joins" or criterion_name is None) | |
| and len(tables) > 1 | |
| and "join" in parsed_query | |
| ): | |
| score += 0.15 | |
| if criterion_name == "aggregation_granularity" or criterion_name is None: | |
| if syntax_res.has_aggregation or any( | |
| agg in parsed_query for agg in ["count(", "sum(", "avg(", "max(", "min("] | |
| ): | |
| score += 0.10 | |
| if syntax_res.has_group_by or "group by" in parsed_query: | |
| score += 0.10 | |
| if ( | |
| (criterion_name == "result_sanity" or criterion_name is None) | |
| and "select" in parsed_query | |
| and not parsed_query.strip().startswith("select *") | |
| ): | |
| score += 0.10 | |
| return round(max(0.0, min(1.0, score)), 4) | |
| def probabilistic_pivot_tournament( | |
| self, | |
| candidates: Sequence[str | dict[str, Any]], | |
| question: str, | |
| schema_ddl: str, | |
| ) -> list[CandidateScore]: | |
| """Rank candidate queries using a Probabilistic Pivot Tournament in O(K * N) time. | |
| Args: | |
| candidates: List or Sequence of SQL strings or dicts with {'sql': str, 'id': str, ...}. | |
| question: Natural language question. | |
| schema_ddl: Database schema DDL string. | |
| Returns: | |
| Ranked list of CandidateScore objects sorted descending by overall_score. | |
| """ | |
| if not candidates: | |
| return [] | |
| scored_candidates: list[CandidateScore] = [] | |
| for idx, cand in enumerate(candidates): | |
| cid = cand.get("id", f"cand_{idx}") if isinstance(cand, dict) else f"cand_{idx}" | |
| sql = cand if isinstance(cand, str) else cand.get("sql", cand.get("query", "")) | |
| exec_valid = cand.get("execution_valid", True) if isinstance(cand, dict) else True | |
| meta = cand if isinstance(cand, dict) else {"original_index": idx} | |
| scored = self.score_query( | |
| sql_query=sql, | |
| question=question, | |
| schema_ddl=schema_ddl, | |
| candidate_id=cid, | |
| execution_valid=exec_valid, | |
| metadata=meta, | |
| ) | |
| scored_candidates.append(scored) | |
| # Sort descending by overall score | |
| scored_candidates.sort(key=lambda c: c.overall_score, reverse=True) | |
| return scored_candidates | |
| def select_best( | |
| self, | |
| candidates: Sequence[str | dict[str, Any]], | |
| question: str, | |
| schema_ddl: str, | |
| ) -> CandidateScore: | |
| """Select the highest scoring candidate via Probabilistic Pivot Tournament.""" | |
| ranked = self.probabilistic_pivot_tournament(candidates, question, schema_ddl) | |
| if not ranked: | |
| raise ValueError("No candidates provided for verification selection.") | |
| return ranked[0] | |