#!/usr/bin/env python3 """ Schema Retrieval Precision/Recall Evaluator for QAFD-RAG Text2SQL. Compares retrieved tables/columns against gold standard. Usage: python benchmarks/text2sql/evaluate_schema.py --results python benchmarks/text2sql/evaluate_schema.py --results --gold """ import os import sys import json import re import argparse from pathlib import Path from typing import Dict, Set, Tuple QAFD_RAG_HOME = Path(__file__).parent.parent.parent DEFAULT_GOLD = QAFD_RAG_HOME / "data" / "text2sql" / "spider2-lite" / "golden_lite_spider_total.json" def load_json_file(file_path: str) -> Dict: with open(file_path, 'r', encoding='utf-8') as f: return json.load(f) def extract_columns_from_gold(data: Dict) -> Dict[str, Set[str]]: """Extract columns from gold standard (CoFD-M format).""" columns_by_instance = {} for instance_id, instance_data in data.items(): if isinstance(instance_data, dict) and 'schema_extraction' in instance_data: schema = instance_data['schema_extraction'] columns_by_instance[instance_id] = set(schema.get('columns', [])) else: columns_by_instance[instance_id] = set() return columns_by_instance def _parse_create_table(create_table_text: str) -> Tuple[Set[str], Set[str]]: """Parse CREATE TABLE statements to extract table and column names.""" tables = set() columns = set() current_table = None for line in create_table_text.splitlines(): stripped = line.strip() table_match = re.match(r'CREATE\s+TABLE\s+`?(\w+)`?\s*\(', stripped, re.IGNORECASE) if table_match: current_table = table_match.group(1) tables.add(current_table) continue if current_table and stripped.startswith('`'): col_match = re.match(r'`(\w+)`\s+\w+', stripped) if col_match: columns.add(f"{current_table}.{col_match.group(1)}") if stripped.startswith(');'): current_table = None return tables, columns def extract_columns_from_results(results_data: Dict) -> Dict[str, Set[str]]: """Extract columns from QAFD-RAG benchmark results by parsing create_table field.""" columns_by_instance = {} for result in results_data.get('results', []): instance_id = result.get('instance_id', '') create_table = result.get('create_table', '') if create_table: _, cols = _parse_create_table(create_table) columns_by_instance[instance_id] = cols else: columns_by_instance[instance_id] = set() return columns_by_instance def extract_tables_from_gold(data: Dict) -> Dict[str, Set[str]]: """Extract tables from gold standard.""" tables_by_instance = {} for instance_id, instance_data in data.items(): if isinstance(instance_data, dict) and 'schema_extraction' in instance_data: schema = instance_data['schema_extraction'] tables_by_instance[instance_id] = set(schema.get('tables', [])) else: tables_by_instance[instance_id] = set() return tables_by_instance def extract_tables_from_results(results_data: Dict) -> Dict[str, Set[str]]: """Extract tables from QAFD-RAG benchmark results by parsing create_table field.""" tables_by_instance = {} for result in results_data.get('results', []): instance_id = result.get('instance_id', '') create_table = result.get('create_table', '') if create_table: tbls, _ = _parse_create_table(create_table) tables_by_instance[instance_id] = tbls else: tables_by_instance[instance_id] = set() return tables_by_instance def calculate_precision_recall(predicted: Set[str], actual: Set[str]) -> Tuple[float, float]: """Calculate precision and recall.""" if not predicted and not actual: return 1.0, 1.0 if not predicted: return 0.0, 0.0 if not actual: return 0.0, 0.0 tp = len(predicted & actual) precision = tp / len(predicted) if predicted else 0.0 recall = tp / len(actual) if actual else 0.0 return precision, recall def evaluate(gold: Dict[str, Set[str]], predicted: Dict[str, Set[str]]) -> Dict: """Evaluate predicted against gold for all common instances.""" common = set(gold.keys()) & set(predicted.keys()) if not common: return {'precision': 0.0, 'recall': 0.0, 'f1': 0.0, 'instances': 0, 'details': {}} precisions, recalls = [], [] details = {} for iid in sorted(common): p, r = calculate_precision_recall(predicted[iid], gold[iid]) precisions.append(p) recalls.append(r) details[iid] = { 'precision': round(p, 4), 'recall': round(r, 4), 'predicted': sorted(predicted[iid]), 'gold': sorted(gold[iid]), 'missing': sorted(gold[iid] - predicted[iid]), 'extra': sorted(predicted[iid] - gold[iid]), } avg_p = sum(precisions) / len(precisions) avg_r = sum(recalls) / len(recalls) f1 = 2 * avg_p * avg_r / (avg_p + avg_r) if (avg_p + avg_r) > 0 else 0.0 return { 'precision': round(avg_p, 4), 'recall': round(avg_r, 4), 'f1': round(f1, 4), 'instances': len(common), 'details': details, } def main(): parser = argparse.ArgumentParser(description="Evaluate text2sql schema retrieval") parser.add_argument("--results", required=True, help="QAFD-RAG benchmark results JSON") parser.add_argument("--gold", default=str(DEFAULT_GOLD), help="Gold standard JSON") parser.add_argument("--output", default=None, help="Save detailed results to JSON") args = parser.parse_args() print("Loading files...") results_data = load_json_file(args.results) gold_data = load_json_file(args.gold) # Extract columns and tables gold_cols = extract_columns_from_gold(gold_data) pred_cols = extract_columns_from_results(results_data) gold_tbls = extract_tables_from_gold(gold_data) pred_tbls = extract_tables_from_results(results_data) # Evaluate col_eval = evaluate(gold_cols, pred_cols) tbl_eval = evaluate(gold_tbls, pred_tbls) # Print results print(f"\n{'=' * 50}") print("SCHEMA RETRIEVAL EVALUATION") print(f"{'=' * 50}") print(f" Instances evaluated: {col_eval['instances']}") print(f"\n COLUMN RETRIEVAL") print(f" {'─' * 40}") print(f" Precision: {col_eval['precision']:.4f}") print(f" Recall: {col_eval['recall']:.4f}") print(f" F1: {col_eval['f1']:.4f}") print(f"\n TABLE RETRIEVAL") print(f" {'─' * 40}") print(f" Precision: {tbl_eval['precision']:.4f}") print(f" Recall: {tbl_eval['recall']:.4f}") print(f" F1: {tbl_eval['f1']:.4f}") # Per-instance details print(f"\n PER-INSTANCE COLUMN DETAILS") print(f" {'─' * 40}") for iid, d in col_eval['details'].items(): status = "OK" if d['recall'] == 1.0 else f"missing: {d['missing']}" print(f" {iid:<15} P={d['precision']:.2f} R={d['recall']:.2f} {status}") # Save output_path = args.output if not output_path: output_path = args.results.replace('.json', '_eval.json') output = { 'column_evaluation': col_eval, 'table_evaluation': tbl_eval, } with open(output_path, 'w') as f: json.dump(output, f, indent=2) print(f"\nDetailed results saved to: {output_path}") if __name__ == "__main__": main()