Add end-to-end evaluation harness for pipeline metrics

scripts/eval_pipeline.py measures per-stage and overall quality:
- Stage 2 retrieval recall@k (fraction of ground-truth tags retrieved)
- Stage 3 selection precision, recall, F1 (final output vs ground truth)
- Per-sample timing for each stage
- Summary with worst/best F1 samples and missed/extra tag analysis

Uses e621_sfw_sample_1000 eval dataset with multiple caption fields.
Supports --skip-rewrite mode and JSONL output for detailed analysis.

https://claude.ai/code/session_019PY5TEXTWGtToUbowunSRG

scripts/eval_pipeline.py

@@ -0,0 +1,391 @@
+"""End-to-end evaluation harness for the Prompt Squirrel RAG pipeline.
+
+Measures per-stage and overall metrics using ground-truth tagged samples
+from the e621 evaluation dataset.
+
+Metrics computed:
+  - Stage 2 (Retrieval): Recall@k — what fraction of ground-truth tags
+    appear among the retrieved candidates
+  - Stage 3 (Selection): Precision, Recall, F1 — how well the final
+    selected tags match the ground truth
+
+Usage:
+    # Full end-to-end (Stage 1 + 2 + 3):
+    python scripts/eval_pipeline.py --n 20
+
+    # Skip Stage 1 LLM rewrite, use ground-truth tags as retrieval input:
+    python scripts/eval_pipeline.py --n 20 --skip-rewrite
+
+    # Use a specific caption field:
+    python scripts/eval_pipeline.py --n 20 --caption-field caption_cogvlm
+
+Requires:
+    - OPENROUTER_API_KEY env var (for Stage 1 rewrite and Stage 3 selection)
+    - fluffyrock_3m.csv and other retrieval assets in the project root
+    - data/eval_samples/e621_sfw_sample_1000_seed123_buffer10000.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Set, Tuple
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+os.chdir(_REPO_ROOT)
+
+EVAL_DATA_PATH = _REPO_ROOT / "data" / "eval_samples" / "e621_sfw_sample_1000_seed123_buffer10000.jsonl"
+
+
+def _flatten_ground_truth_tags(tags_categorized_str: str) -> Set[str]:
+    """Parse the categorized ground-truth JSON string into a flat set of tags."""
+    if not tags_categorized_str:
+        return set()
+    try:
+        cats = json.loads(tags_categorized_str)
+    except json.JSONDecodeError:
+        return set()
+    tags = set()
+    for tag_list in cats.values():
+        if isinstance(tag_list, list):
+            for t in tag_list:
+                tags.add(t.strip())
+    return tags
+
+
+@dataclass
+class SampleResult:
+    sample_id: Any
+    caption: str
+    ground_truth_tags: Set[str]
+    # Stage 1
+    rewrite_phrases: List[str] = field(default_factory=list)
+    # Stage 2
+    retrieved_tags: Set[str] = field(default_factory=set)
+    retrieval_recall: float = 0.0
+    # Stage 3
+    selected_tags: Set[str] = field(default_factory=set)
+    selection_precision: float = 0.0
+    selection_recall: float = 0.0
+    selection_f1: float = 0.0
+    # Timing
+    stage1_time: float = 0.0
+    stage2_time: float = 0.0
+    stage3_time: float = 0.0
+    # Errors
+    error: Optional[str] = None
+
+
+def _compute_metrics(predicted: Set[str], ground_truth: Set[str]) -> Tuple[float, float, float]:
+    """Compute precision, recall, F1."""
+    if not predicted and not ground_truth:
+        return 1.0, 1.0, 1.0
+    if not predicted:
+        return 0.0, 0.0, 0.0
+    if not ground_truth:
+        return 0.0, 0.0, 0.0
+
+    tp = len(predicted & ground_truth)
+    precision = tp / len(predicted)
+    recall = tp / len(ground_truth)
+    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
+    return precision, recall, f1
+
+
+def run_eval(
+    n_samples: int = 20,
+    caption_field: str = "caption_cogvlm",
+    skip_rewrite: bool = False,
+    allow_nsfw: bool = False,
+    mode: str = "chunked_map_union",
+    chunk_size: int = 60,
+    per_phrase_k: int = 2,
+    temperature: float = 0.0,
+    max_tokens: int = 512,
+    verbose: bool = False,
+) -> List[SampleResult]:
+
+    from psq_rag.llm.rewrite import llm_rewrite_prompt
+    from psq_rag.retrieval.psq_retrieval import psq_candidates_from_rewrite_phrases
+    from psq_rag.llm.select import llm_select_indices
+
+    def log(msg: str) -> None:
+        if verbose:
+            print(f"  {msg}")
+
+    # Load eval samples
+    if not EVAL_DATA_PATH.is_file():
+        print(f"ERROR: Eval data not found: {EVAL_DATA_PATH}")
+        sys.exit(1)
+
+    samples = []
+    with EVAL_DATA_PATH.open("r", encoding="utf-8") as f:
+        for line in f:
+            if len(samples) >= n_samples:
+                break
+            row = json.loads(line)
+            caption = row.get(caption_field, "")
+            if not caption or not caption.strip():
+                continue
+            gt_tags = _flatten_ground_truth_tags(row.get("tags_ground_truth_categorized", ""))
+            if not gt_tags:
+                continue
+            samples.append({
+                "id": row.get("id", row.get("row_id", len(samples))),
+                "caption": caption.strip(),
+                "gt_tags": gt_tags,
+            })
+
+    print(f"Loaded {len(samples)} samples (caption_field={caption_field})")
+    print(f"skip_rewrite={skip_rewrite}, allow_nsfw={allow_nsfw}, mode={mode}")
+    print()
+
+    results: List[SampleResult] = []
+
+    for i, sample in enumerate(samples):
+        sid = sample["id"]
+        caption = sample["caption"]
+        gt_tags = sample["gt_tags"]
+
+        result = SampleResult(
+            sample_id=sid,
+            caption=caption[:120] + ("..." if len(caption) > 120 else ""),
+            ground_truth_tags=gt_tags,
+        )
+
+        print(f"[{i+1}/{len(samples)}] id={sid} gt_tags={len(gt_tags)}")
+
+        try:
+            # --- Stage 1: LLM Rewrite ---
+            if skip_rewrite:
+                # Use the caption directly as comma-separated phrases
+                phrases = [p.strip() for p in caption.split(",") if p.strip()]
+                # Also split on periods/sentences for natural language captions
+                if len(phrases) <= 1:
+                    phrases = [p.strip() for p in caption.replace(".", ",").split(",") if p.strip()]
+                result.rewrite_phrases = phrases
+                result.stage1_time = 0.0
+            else:
+                t0 = time.time()
+                rewritten = llm_rewrite_prompt(caption, log)
+                result.stage1_time = time.time() - t0
+                if rewritten:
+                    result.rewrite_phrases = [p.strip() for p in rewritten.split(",") if p.strip()]
+                else:
+                    result.rewrite_phrases = [p.strip() for p in caption.split(",") if p.strip()]
+                    if len(result.rewrite_phrases) <= 1:
+                        result.rewrite_phrases = [p.strip() for p in caption.replace(".", ",").split(",") if p.strip()]
+
+            if verbose:
+                log(f"Phrases ({len(result.rewrite_phrases)}): {result.rewrite_phrases[:5]}")
+
+            # --- Stage 2: Retrieval ---
+            t0 = time.time()
+            retrieval_result = psq_candidates_from_rewrite_phrases(
+                rewrite_phrases=result.rewrite_phrases,
+                allow_nsfw_tags=allow_nsfw,
+                global_k=300,
+                verbose=False,
+            )
+            result.stage2_time = time.time() - t0
+
+            if isinstance(retrieval_result, tuple):
+                candidates, _ = retrieval_result
+            else:
+                candidates = retrieval_result
+
+            result.retrieved_tags = {c.tag for c in candidates}
+            # Retrieval recall: what fraction of ground truth was retrieved
+            if gt_tags:
+                result.retrieval_recall = len(result.retrieved_tags & gt_tags) / len(gt_tags)
+
+            if verbose:
+                log(f"Retrieved {len(candidates)} candidates, recall={result.retrieval_recall:.3f}")
+
+            # --- Stage 3: LLM Selection ---
+            t0 = time.time()
+            picked_indices = llm_select_indices(
+                query_text=caption,
+                candidates=candidates,
+                max_pick=0,
+                log=log,
+                mode=mode,
+                chunk_size=chunk_size,
+                per_phrase_k=per_phrase_k,
+                temperature=temperature,
+                max_tokens=max_tokens,
+            )
+            result.stage3_time = time.time() - t0
+
+            result.selected_tags = {candidates[idx].tag for idx in picked_indices} if picked_indices else set()
+
+            # Selection metrics
+            p, r, f1 = _compute_metrics(result.selected_tags, gt_tags)
+            result.selection_precision = p
+            result.selection_recall = r
+            result.selection_f1 = f1
+
+            print(
+                f"  retrieval_recall={result.retrieval_recall:.3f} "
+                f"sel_P={p:.3f} sel_R={r:.3f} sel_F1={f1:.3f} "
+                f"selected={len(result.selected_tags)} "
+                f"t1={result.stage1_time:.1f}s t2={result.stage2_time:.1f}s t3={result.stage3_time:.1f}s"
+            )
+
+        except Exception as e:
+            result.error = str(e)
+            print(f"  ERROR: {e}")
+
+        results.append(result)
+
+    return results
+
+
+def print_summary(results: List[SampleResult]) -> None:
+    """Print aggregate metrics across all samples."""
+    valid = [r for r in results if r.error is None]
+    errored = [r for r in results if r.error is not None]
+
+    if not valid:
+        print("\nNo valid results to summarize.")
+        return
+
+    n = len(valid)
+
+    avg_retrieval_recall = sum(r.retrieval_recall for r in valid) / n
+    avg_sel_precision = sum(r.selection_precision for r in valid) / n
+    avg_sel_recall = sum(r.selection_recall for r in valid) / n
+    avg_sel_f1 = sum(r.selection_f1 for r in valid) / n
+
+    avg_retrieved = sum(len(r.retrieved_tags) for r in valid) / n
+    avg_selected = sum(len(r.selected_tags) for r in valid) / n
+    avg_gt = sum(len(r.ground_truth_tags) for r in valid) / n
+
+    avg_t1 = sum(r.stage1_time for r in valid) / n
+    avg_t2 = sum(r.stage2_time for r in valid) / n
+    avg_t3 = sum(r.stage3_time for r in valid) / n
+
+    print()
+    print("=" * 60)
+    print(f"EVALUATION SUMMARY ({n} samples, {len(errored)} errors)")
+    print("=" * 60)
+    print()
+    print("Stage 2 - Retrieval:")
+    print(f"  Avg recall@300:       {avg_retrieval_recall:.4f}")
+    print(f"  Avg candidates:       {avg_retrieved:.1f}")
+    print()
+    print("Stage 3 - Selection (final output):")
+    print(f"  Avg precision:        {avg_sel_precision:.4f}")
+    print(f"  Avg recall:           {avg_sel_recall:.4f}")
+    print(f"  Avg F1:               {avg_sel_f1:.4f}")
+    print(f"  Avg selected tags:    {avg_selected:.1f}")
+    print(f"  Avg ground-truth tags:{avg_gt:.1f}")
+    print()
+    print("Timing (avg per sample):")
+    print(f"  Stage 1 (rewrite):    {avg_t1:.2f}s")
+    print(f"  Stage 2 (retrieval):  {avg_t2:.2f}s")
+    print(f"  Stage 3 (selection):  {avg_t3:.2f}s")
+    print(f"  Total:                {avg_t1 + avg_t2 + avg_t3:.2f}s")
+    print()
+
+    # Show worst and best F1 samples
+    by_f1 = sorted(valid, key=lambda r: r.selection_f1)
+    print("Lowest F1 samples:")
+    for r in by_f1[:3]:
+        print(f"  id={r.sample_id} F1={r.selection_f1:.3f} P={r.selection_precision:.3f} R={r.selection_recall:.3f}")
+        missed = r.ground_truth_tags - r.selected_tags
+        extra = r.selected_tags - r.ground_truth_tags
+        if missed:
+            print(f"    missed: {sorted(missed)[:10]}")
+        if extra:
+            print(f"    extra:  {sorted(extra)[:10]}")
+
+    print()
+    print("Highest F1 samples:")
+    for r in by_f1[-3:]:
+        print(f"  id={r.sample_id} F1={r.selection_f1:.3f} P={r.selection_precision:.3f} R={r.selection_recall:.3f}")
+
+    if errored:
+        print()
+        print(f"Errors ({len(errored)}):")
+        for r in errored[:5]:
+            print(f"  id={r.sample_id}: {r.error}")
+
+    print("=" * 60)
+
+
+def main(argv=None) -> int:
+    ap = argparse.ArgumentParser(description="End-to-end pipeline evaluation")
+    ap.add_argument("--n", type=int, default=20, help="Number of samples to evaluate")
+    ap.add_argument("--caption-field", default="caption_cogvlm",
+                    choices=["caption_cogvlm", "caption_llm_0", "caption_llm_1",
+                             "caption_llm_2", "caption_llm_3", "caption_llm_4",
+                             "caption_llm_5", "caption_llm_6", "caption_llm_7"],
+                    help="Which caption field to use as input")
+    ap.add_argument("--skip-rewrite", action="store_true",
+                    help="Skip Stage 1 LLM rewrite; split caption directly into phrases")
+    ap.add_argument("--allow-nsfw", action="store_true", help="Allow NSFW tags")
+    ap.add_argument("--mode", default="chunked_map_union",
+                    choices=["single_shot", "chunked_map_union"])
+    ap.add_argument("--chunk-size", type=int, default=60)
+    ap.add_argument("--per-phrase-k", type=int, default=2)
+    ap.add_argument("--temperature", type=float, default=0.0)
+    ap.add_argument("--max-tokens", type=int, default=512)
+    ap.add_argument("--verbose", "-v", action="store_true", help="Show per-call Stage 3 logs")
+    ap.add_argument("--output", "-o", type=str, default=None,
+                    help="Save detailed results as JSONL to this path")
+
+    args = ap.parse_args(list(argv) if argv is not None else None)
+
+    results = run_eval(
+        n_samples=args.n,
+        caption_field=args.caption_field,
+        skip_rewrite=args.skip_rewrite,
+        allow_nsfw=args.allow_nsfw,
+        mode=args.mode,
+        chunk_size=args.chunk_size,
+        per_phrase_k=args.per_phrase_k,
+        temperature=args.temperature,
+        max_tokens=args.max_tokens,
+        verbose=args.verbose,
+    )
+
+    print_summary(results)
+
+    # Optionally save detailed results
+    if args.output:
+        out_path = Path(args.output)
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        with out_path.open("w", encoding="utf-8") as f:
+            for r in results:
+                row = {
+                    "sample_id": r.sample_id,
+                    "caption": r.caption,
+                    "ground_truth_tags": sorted(r.ground_truth_tags),
+                    "rewrite_phrases": r.rewrite_phrases,
+                    "retrieved_tags": sorted(r.retrieved_tags),
+                    "selected_tags": sorted(r.selected_tags),
+                    "retrieval_recall": round(r.retrieval_recall, 4),
+                    "selection_precision": round(r.selection_precision, 4),
+                    "selection_recall": round(r.selection_recall, 4),
+                    "selection_f1": round(r.selection_f1, 4),
+                    "stage1_time": round(r.stage1_time, 3),
+                    "stage2_time": round(r.stage2_time, 3),
+                    "stage3_time": round(r.stage3_time, 3),
+                    "error": r.error,
+                }
+                f.write(json.dumps(row, ensure_ascii=False) + "\n")
+        print(f"\nDetailed results saved to: {out_path}")
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())