# /// script

# requires-python = ">=3.10"

# dependencies = [

# "transformers>=4.45.0",

# "peft>=0.7.0",

# "torch",

# "accelerate",

# ]

# ///

"""

QMD Retrieval-Based Evaluation with Precision & Recall

Evaluates model outputs against golden data (training set).

Measures how well the model reproduces the expected expansions.

Metrics:

- Precision: Of model-generated expansions, how many match golden?

- Recall: Of golden expansions, how many did the model generate?

- F1: Harmonic mean of precision and recall

Matching is done via token overlap (Jaccard similarity) with a threshold.

Usage:

uv run eval_retrieval.py ./outputs/sft

uv run eval_retrieval.py tobil/qmd-query-expansion-1.7B --golden data/qmd_expansion_v3_structured.jsonl

uv run eval_retrieval.py ./outputs/sft --threshold 0.5 --sample 100

"""

import argparse

import json

import random

import re

import sys

from collections import defaultdict

from pathlib import Path

# =============================================================================

# Matching Functions

# =============================================================================

def tokenize(text: str) -> set[str]:

"""Tokenize text into lowercase word set, removing stopwords."""

stopwords = {'the', 'a', 'an', 'is', 'are', 'to', 'for', 'of', 'in', 'and',

'or', 'it', 'this', 'that', 'be', 'with', 'as', 'on', 'by',

'how', 'what', 'do', 'does', 'can', 'you', 'your', 'i'}

words = re.findall(r'\b\w+\b', text.lower())

return {w for w in words if w not in stopwords and len(w) > 1}

def jaccard_similarity(a: str, b: str) -> float:

"""Jaccard similarity between two strings based on token overlap."""

tokens_a = tokenize(a)

tokens_b = tokenize(b)

if not tokens_a or not tokens_b:

return 0.0

intersection = len(tokens_a & tokens_b)

union = len(tokens_a | tokens_b)

return intersection / union if union > 0 else 0.0

def find_best_match(pred: str, golden_list: list[str], threshold: float) -> tuple[str | None, float]:

"""Find best matching golden expansion for a prediction."""

best_match = None

best_score = 0.0

for golden in golden_list:

score = jaccard_similarity(pred, golden)

if score > best_score:

best_score = score

best_match = golden

if best_score >= threshold:

return best_match, best_score

return None, best_score

# =============================================================================

# Parsing

# =============================================================================

def parse_model_output(text: str) -> dict[str, list[str]]:

"""Parse model output into {lex: [...], vec: [...], hyde: [...]}."""

# Clean thinking tags

text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)

text = text.replace('<|im_end|>', '').strip()

result = {"lex": [], "vec": [], "hyde": []}

for line in text.strip().split("\n"):

line = line.strip()

if not line:

continue

if line.startswith("lex:"):

result["lex"].append(line[4:].strip())

elif line.startswith("vec:"):

result["vec"].append(line[4:].strip())

elif line.startswith("hyde:"):

result["hyde"].append(line[5:].strip())

return result

def parse_golden_data(searches: list[dict] | str) -> dict[str, list[str]]:

"""Parse golden data format into {lex: [...], vec: [...], hyde: [...]}."""

# If it's a string (from messages format), parse it

if isinstance(searches, str):

return parse_model_output(searches)

# Otherwise it's the structured format [{type, query}, ...]

result = {"lex": [], "vec": [], "hyde": []}

for item in searches:

exp_type = item.get("type", "")

value = item.get("query", "") or item.get("value", "")

if exp_type in result:

result[exp_type].append(value)

return result

def load_golden_data(filepath: Path) -> list[dict]:

"""Load golden data from JSONL, supporting both structured and messages formats."""

data = []

with open(filepath) as f:

for line in f:

if not line.strip():

continue

item = json.loads(line)

# Structured format: {query, searches}

if "query" in item and "searches" in item:

data.append({

"query": item["query"],

"searches": item["searches"]

})

# Messages format: {messages: [{role, content}, ...]}

elif "messages" in item:

messages = item["messages"]

query = None

searches = None

for msg in messages:

if msg["role"] == "user":

# Extract query from "/no_think Expand this search query: ..."

content = msg["content"]

if "Expand this search query:" in content:

query = content.split("Expand this search query:")[-1].strip()

else:

query = content.strip()

elif msg["role"] == "assistant":

# The assistant content IS the expected output

searches = msg["content"]

if query and searches:

data.append({

"query": query,

"searches": searches # Will be parsed as string

})

return data

# =============================================================================

# Metrics Calculation

# =============================================================================

# Different thresholds by type - lex needs strict matching, hyde is more flexible

DEFAULT_THRESHOLDS = {

"lex": 0.5, # Keywords should overlap well

"vec": 0.35, # Semantic sentences have more variation

"hyde": 0.25, # Passages have the most variation

}

def calculate_metrics(

predictions: dict[str, list[str]],

golden: dict[str, list[str]],

threshold: float | dict[str, float] = 0.4,

return_mismatches: bool = False

) -> dict:

"""Calculate precision, recall, F1 per type and overall.

Args:

threshold: Either a single float, or dict mapping type -> threshold

return_mismatches: If True, include lists of unmatched predictions/golden

"""

if isinstance(threshold, (int, float)):

thresholds = {"lex": threshold, "vec": threshold, "hyde": threshold}

else:

thresholds = threshold

metrics = {}

mismatches = {}

total_tp = 0

total_pred = 0

total_golden = 0

for exp_type in ["lex", "vec", "hyde"]:

preds = predictions.get(exp_type, [])

golds = golden.get(exp_type, [])

type_threshold = thresholds.get(exp_type, 0.4)

if not preds and not golds:

continue

# Track which golden items were matched

matched_golden = set()

unmatched_preds = []

tp = 0

for pred in preds:

match, score = find_best_match(pred, golds, type_threshold)

if match is not None:

tp += 1

matched_golden.add(match)

else:

unmatched_preds.append((pred, score))

unmatched_golden = [g for g in golds if g not in matched_golden]

precision = tp / len(preds) if preds else 0.0

recall = len(matched_golden) / len(golds) if golds else 0.0

f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0

metrics[exp_type] = {

"precision": precision,

"recall": recall,

"f1": f1,

"pred_count": len(preds),

"golden_count": len(golds),

"matched": tp,

}

if return_mismatches:

mismatches[exp_type] = {

"unmatched_preds": unmatched_preds,

"unmatched_golden": unmatched_golden,

}

total_tp += tp

total_pred += len(preds)

total_golden += len(golds)

# Overall metrics (micro-averaged)

overall_precision = total_tp / total_pred if total_pred > 0 else 0.0

overall_recall = total_tp / total_golden if total_golden > 0 else 0.0

overall_f1 = 2 * overall_precision * overall_recall / (overall_precision + overall_recall) if (overall_precision + overall_recall) > 0 else 0.0

metrics["overall"] = {

"precision": overall_precision,

"recall": overall_recall,

"f1": overall_f1,

"pred_count": total_pred,

"golden_count": total_golden,

"matched": total_tp,

}

if return_mismatches:

metrics["_mismatches"] = mismatches

return metrics

# =============================================================================

# Model Loading and Generation

# =============================================================================

def load_model(model_path: str):

"""Load model (adapter or merged)."""

import torch

from peft import PeftModel

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

model_path = Path(model_path)

adapter_config = model_path / "adapter_config.json"

# Get base model from adapter config or default

base_model = "Qwen/Qwen3-1.7B"

if adapter_config.exists():

with open(adapter_config) as f:

cfg = json.load(f)

base_model = cfg.get("base_model_name_or_path", base_model)

print(f"Loading base: {base_model}", file=sys.stderr)

tokenizer = AutoTokenizer.from_pretrained(base_model)

if tokenizer.pad_token is None:

tokenizer.pad_token = tokenizer.eos_token

tokenizer.padding_side = "left"

config = AutoConfig.from_pretrained(base_model)

config.tie_word_embeddings = False

model = AutoModelForCausalLM.from_pretrained(

base_model, dtype=torch.bfloat16, device_map={"": 0}, config=config

)

if model.generation_config is not None:

model.generation_config.do_sample = False

model.generation_config.temperature = None

model.generation_config.top_p = None

model.generation_config.top_k = None

# Load adapter if present

if adapter_config.exists():

print(f"Loading adapter: {model_path}", file=sys.stderr)

model = PeftModel.from_pretrained(model, str(model_path))

model.eval()

return model, tokenizer

def generate_expansion(model, tokenizer, query: str, max_new_tokens: int = 400) -> str:

"""Generate expansion for a single query."""

import torch

prompt = tokenizer.apply_chat_template(

[{"role": "user", "content": f"/no_think Expand this search query: {query}"}],

tokenize=False,

add_generation_prompt=True,

)

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

input_len = inputs["input_ids"].shape[1]

with torch.inference_mode():

out = model.generate(

**inputs,

max_new_tokens=max_new_tokens,

do_sample=False,

num_beams=1,

pad_token_id=tokenizer.pad_token_id,

eos_token_id=tokenizer.eos_token_id,

use_cache=True,

)

gen_tokens = out[0][input_len:]

return tokenizer.decode(gen_tokens, skip_special_tokens=True)

# =============================================================================

# Main Evaluation

# =============================================================================

def main():

parser = argparse.ArgumentParser(description="QMD Retrieval-Based Evaluation")

parser.add_argument("model", help="Model path (local or HF)")

parser.add_argument("--golden", default="data/qmd_expansion_v3_structured.jsonl",

help="Golden data JSONL file")

parser.add_argument("--threshold", type=float, default=None,

help="Jaccard similarity threshold for all types (overrides --type-thresholds)")

parser.add_argument("--type-thresholds", action="store_true",

help="Use type-specific thresholds (lex=0.5, vec=0.35, hyde=0.25)")

parser.add_argument("--sample", type=int, default=0,

help="Sample N queries (0 = all)")

parser.add_argument("--seed", type=int, default=42,

help="Random seed for sampling")

parser.add_argument("--max-new-tokens", type=int, default=400,

help="Max new tokens to generate")

parser.add_argument("--verbose", "-v", action="store_true",

help="Show per-query details")

parser.add_argument("--show-mismatches", action="store_true",

help="Show examples of mismatched predictions")

args = parser.parse_args()

# Determine thresholds

if args.threshold is not None:

thresholds = args.threshold

elif args.type_thresholds:

thresholds = DEFAULT_THRESHOLDS.copy()

else:

thresholds = 0.4 # Default single threshold

# Load golden data

golden_path = Path(args.golden)

if not golden_path.exists():

# Try relative to script directory

golden_path = Path(__file__).parent / args.golden

if not golden_path.exists():

print(f"Error: Golden data file not found: {args.golden}", file=sys.stderr)

sys.exit(1)

print(f"Loading golden data from {golden_path}...", file=sys.stderr)

golden_data = load_golden_data(golden_path)

print(f"Loaded {len(golden_data)} golden examples", file=sys.stderr)

# Sample if requested

if args.sample > 0 and args.sample < len(golden_data):

random.seed(args.seed)

golden_data = random.sample(golden_data, args.sample)

print(f"Sampled {len(golden_data)} examples", file=sys.stderr)

# Load model

model, tokenizer = load_model(args.model)

# Evaluate

all_metrics = []

all_mismatches = []

type_aggregates = defaultdict(lambda: {"precision": [], "recall": [], "f1": []})

threshold_desc = thresholds if isinstance(thresholds, (int, float)) else f"lex={thresholds['lex']}, vec={thresholds['vec']}, hyde={thresholds['hyde']}"

print(f"\nEvaluating {len(golden_data)} queries (thresholds: {threshold_desc})...\n")

for i, item in enumerate(golden_data, 1):

query = item["query"]

golden_parsed = parse_golden_data(item["searches"])

# Generate model output

output = generate_expansion(model, tokenizer, query, args.max_new_tokens)

pred_parsed = parse_model_output(output)

# Calculate metrics

metrics = calculate_metrics(pred_parsed, golden_parsed, thresholds, return_mismatches=args.show_mismatches)

all_metrics.append({"query": query, "metrics": metrics, "pred": pred_parsed, "golden": golden_parsed})

if args.show_mismatches and "_mismatches" in metrics:

all_mismatches.append({"query": query, "mismatches": metrics.pop("_mismatches")})

# Aggregate by type

for exp_type in ["lex", "vec", "hyde", "overall"]:

if exp_type in metrics:

type_aggregates[exp_type]["precision"].append(metrics[exp_type]["precision"])

type_aggregates[exp_type]["recall"].append(metrics[exp_type]["recall"])

type_aggregates[exp_type]["f1"].append(metrics[exp_type]["f1"])

# Progress

overall = metrics.get("overall", {})

p = overall.get("precision", 0) * 100

r = overall.get("recall", 0) * 100

f = overall.get("f1", 0) * 100

if args.verbose:

print(f"[{i:3d}/{len(golden_data)}] P={p:5.1f}% R={r:5.1f}% F1={f:5.1f}% {query[:50]}")

elif i % 50 == 0 or i == len(golden_data):

print(f" Processed {i}/{len(golden_data)}...", file=sys.stderr)

# Summary

print(f"\n{'='*60}")

print(f"RESULTS: {args.model}")

print(f"{'='*60}")

print(f"Threshold: {args.threshold} | Samples: {len(golden_data)}")

print()

print(f"{'Type':<10} {'Precision':>10} {'Recall':>10} {'F1':>10}")

print("-" * 42)

for exp_type in ["lex", "vec", "hyde", "overall"]:

if exp_type in type_aggregates:

agg = type_aggregates[exp_type]

avg_p = sum(agg["precision"]) / len(agg["precision"]) * 100 if agg["precision"] else 0

avg_r = sum(agg["recall"]) / len(agg["recall"]) * 100 if agg["recall"] else 0

avg_f = sum(agg["f1"]) / len(agg["f1"]) * 100 if agg["f1"] else 0

label = exp_type.upper() if exp_type != "overall" else "OVERALL"

print(f"{label:<10} {avg_p:>9.1f}% {avg_r:>9.1f}% {avg_f:>9.1f}%")

print(f"{'='*60}")

# Show worst examples

print("\nBottom 5 by F1:")

sorted_by_f1 = sorted(all_metrics, key=lambda x: x["metrics"].get("overall", {}).get("f1", 0))

for item in sorted_by_f1[:5]:

f1 = item["metrics"].get("overall", {}).get("f1", 0) * 100

print(f" {f1:5.1f}% {item['query'][:60]}")

# Show mismatches if requested

if args.show_mismatches and all_mismatches:

print(f"\n{'='*60}")

print("MISMATCH EXAMPLES")

print(f"{'='*60}")

# Group by type and show up to 3 examples per type

for exp_type in ["lex", "vec", "hyde"]:

type_mismatches = []

for item in all_mismatches:

if exp_type in item["mismatches"]:

mm = item["mismatches"][exp_type]

if mm["unmatched_preds"] or mm["unmatched_golden"]:

type_mismatches.append({

"query": item["query"],

**mm

})

if type_mismatches:

print(f"\n--- {exp_type.upper()} mismatches ({len(type_mismatches)} queries) ---")

for example in type_mismatches[:3]:

print(f"\nQuery: {example['query'][:60]}")

if example["unmatched_preds"]:

print(f" Unmatched predictions:")

for pred, score in example["unmatched_preds"][:2]:

print(f" - [{score:.2f}] {pred[:80]}{'...' if len(pred) > 80 else ''}")

if example["unmatched_golden"]:

print(f" Missing golden:")

for g in example["unmatched_golden"][:2]:

print(f" - {g[:80]}{'...' if len(g) > 80 else ''}")

return 0

if __name__ == "__main__":

sys.exit(main())