Add parallel processing to eval pipeline with ThreadPoolExecutor

Samples are now processed concurrently (default 4 workers) using threads,
which is ideal since the bottleneck is I/O (OpenRouter API calls). Retrieval
assets are pre-warmed before threads start to avoid initialization races.
Use --workers 1 to revert to sequential mode.

https://claude.ai/code/session_019PY5TEXTWGtToUbowunSRG

scripts/eval_pipeline.py

@@ -16,6 +16,12 @@ Usage:
     # Reproducible run with specific seed:
     python scripts/eval_pipeline.py --n 50 --seed 123
 
+    # Parallel processing with 4 workers (default):
+    python scripts/eval_pipeline.py --n 50 --workers 4
+
+    # Sequential mode (disable parallelism):
+    python scripts/eval_pipeline.py --n 20 --workers 1
+
     # Skip Stage 1 LLM rewrite (cheaper, tests Stage 2+3 only):
     python scripts/eval_pipeline.py --n 20 --skip-rewrite
 
@@ -38,7 +44,9 @@ import json
 import os
 import random
 import sys
+import threading
 import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
@@ -143,6 +151,169 @@ def _compute_metrics(predicted: Set[str], ground_truth: Set[str]) -> Tuple[float
     return precision, recall, f1
 
 
+def _process_one_sample(
+    sample: Dict[str, Any],
+    index: int,
+    total: int,
+    skip_rewrite: bool,
+    allow_nsfw: bool,
+    mode: str,
+    chunk_size: int,
+    per_phrase_k: int,
+    temperature: float,
+    max_tokens: int,
+    verbose: bool,
+    print_lock: threading.Lock,
+) -> SampleResult:
+    """Process a single eval sample through the full pipeline. Thread-safe."""
+    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
+    from psq_rag.retrieval.state import get_tag_type_name
+
+    def log(msg: str) -> None:
+        if verbose:
+            with print_lock:
+                print(f"  [{index+1}] {msg}")
+
+    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,
+    )
+
+    with print_lock:
+        print(f"[{index+1}/{total}] id={sid} gt_tags={len(gt_tags)}")
+
+    try:
+        # --- Stage 1: LLM Rewrite ---
+        if skip_rewrite:
+            phrases = [p.strip() for p in caption.split(",") if p.strip()]
+            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()]
+
+        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}
+        if gt_tags:
+            result.retrieval_recall = len(result.retrieved_tags & gt_tags) / len(gt_tags)
+
+        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()
+
+        # Overall selection metrics
+        p, r, f1 = _compute_metrics(result.selected_tags, gt_tags)
+        result.selection_precision = p
+        result.selection_recall = r
+        result.selection_f1 = f1
+
+        # Split ground-truth and selected tags by type
+        gt_char, gt_gen = _classify_tags(gt_tags, get_tag_type_name)
+        sel_char, sel_gen = _classify_tags(result.selected_tags, get_tag_type_name)
+        ret_char, _ = _classify_tags(result.retrieved_tags, get_tag_type_name)
+
+        result.gt_character_tags = gt_char
+        result.selected_character_tags = sel_char
+        result.retrieved_character_tags = ret_char
+        result.gt_general_tags = gt_gen
+        result.selected_general_tags = sel_gen
+
+        # Character-specific metrics
+        if gt_char:
+            result.char_retrieval_recall = len(ret_char & gt_char) / len(gt_char)
+        cp, cr, cf1 = _compute_metrics(sel_char, gt_char)
+        result.char_precision = cp
+        result.char_recall = cr
+        result.char_f1 = cf1
+
+        # General-tag metrics
+        gp, gr, gf1 = _compute_metrics(sel_gen, gt_gen)
+        result.general_precision = gp
+        result.general_recall = gr
+        result.general_f1 = gf1
+
+        # Per-sample output line
+        char_info = ""
+        if gt_char:
+            char_info = f" char[gt={len(gt_char)} sel={len(sel_char)} P={cp:.2f} R={cr:.2f}]"
+        with print_lock:
+            print(
+                f"  [{index+1}] 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)}{char_info} "
+                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)
+        with print_lock:
+            print(f"  [{index+1}] ERROR: {e}")
+
+    return result
+
+
+def _prewarm_retrieval_assets() -> None:
+    """Force-load all lazy retrieval assets so threads don't race on init."""
+    from psq_rag.retrieval.state import (
+        get_tfidf_components,
+        get_tag2aliases,
+        get_tag_type_name,
+    )
+    print("Pre-warming retrieval assets (TF-IDF, FastText, HNSW, aliases)...")
+    t0 = time.time()
+    get_tfidf_components()  # loads joblib, HNSW indexes, FastText model
+    get_tag2aliases()       # loads CSV alias dict
+    get_tag_type_name("_warmup_")  # ensures tag type dict is built
+    print(f"  Assets loaded in {time.time() - t0:.1f}s")
+
+
 def run_eval(
     n_samples: int = 20,
     caption_field: str = "caption_cogvlm",
@@ -156,17 +327,9 @@ def run_eval(
     verbose: bool = False,
     shuffle: bool = True,
     seed: int = 42,
+    workers: int = 1,
 ) -> 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
-    from psq_rag.retrieval.state import get_tag_type_name
-
-    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}")
@@ -196,134 +359,63 @@ def run_eval(
 
     print(f"Loaded {len(samples)}/{len(all_samples)} samples (caption_field={caption_field})")
     print(f"shuffle={shuffle}, seed={seed}, skip_rewrite={skip_rewrite}, allow_nsfw={allow_nsfw}, mode={mode}")
+    print(f"workers={workers}")
     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}")
+    # Pre-warm shared retrieval assets before spawning threads
+    _prewarm_retrieval_assets()
+    print()
 
-            # --- 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,
+    print_lock = threading.Lock()
+    total = len(samples)
+
+    if workers <= 1:
+        # Sequential mode (original behavior)
+        results: List[SampleResult] = []
+        for i, sample in enumerate(samples):
+            result = _process_one_sample(
+                sample, i, total,
+                skip_rewrite, allow_nsfw, mode, chunk_size,
+                per_phrase_k, temperature, max_tokens, verbose,
+                print_lock,
             )
-            result.stage3_time = time.time() - t0
-
-            result.selected_tags = {candidates[idx].tag for idx in picked_indices} if picked_indices else set()
-
-            # Overall selection metrics
-            p, r, f1 = _compute_metrics(result.selected_tags, gt_tags)
-            result.selection_precision = p
-            result.selection_recall = r
-            result.selection_f1 = f1
-
-            # Split ground-truth and selected tags by type
-            gt_char, gt_gen = _classify_tags(gt_tags, get_tag_type_name)
-            sel_char, sel_gen = _classify_tags(result.selected_tags, get_tag_type_name)
-            ret_char, _ = _classify_tags(result.retrieved_tags, get_tag_type_name)
-
-            result.gt_character_tags = gt_char
-            result.selected_character_tags = sel_char
-            result.retrieved_character_tags = ret_char
-            result.gt_general_tags = gt_gen
-            result.selected_general_tags = sel_gen
-
-            # Character-specific metrics
-            if gt_char:
-                result.char_retrieval_recall = len(ret_char & gt_char) / len(gt_char)
-            cp, cr, cf1 = _compute_metrics(sel_char, gt_char)
-            result.char_precision = cp
-            result.char_recall = cr
-            result.char_f1 = cf1
-
-            # General-tag metrics
-            gp, gr, gf1 = _compute_metrics(sel_gen, gt_gen)
-            result.general_precision = gp
-            result.general_recall = gr
-            result.general_f1 = gf1
-
-            # Per-sample output line
-            char_info = ""
-            if gt_char:
-                char_info = f" char[gt={len(gt_char)} sel={len(sel_char)} P={cp:.2f} R={cr:.2f}]"
-            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)}{char_info} "
-                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)
+            results.append(result)
+    else:
+        # Parallel mode
+        print(f"Processing {total} samples with {workers} parallel workers...")
+        print()
+        # Submit all samples; use index to preserve original ordering
+        results_by_index: Dict[int, SampleResult] = {}
+        with ThreadPoolExecutor(max_workers=workers) as executor:
+            futures = {
+                executor.submit(
+                    _process_one_sample,
+                    sample, i, total,
+                    skip_rewrite, allow_nsfw, mode, chunk_size,
+                    per_phrase_k, temperature, max_tokens, verbose,
+                    print_lock,
+                ): i
+                for i, sample in enumerate(samples)
+            }
+            for future in as_completed(futures):
+                idx = futures[future]
+                try:
+                    results_by_index[idx] = future.result()
+                except Exception as e:
+                    # Should not happen since _process_one_sample catches exceptions,
+                    # but guard against unexpected errors
+                    with print_lock:
+                        print(f"  [{idx+1}] WORKER ERROR: {e}")
+                    result = SampleResult(
+                        sample_id=samples[idx]["id"],
+                        caption=samples[idx]["caption"][:120],
+                        ground_truth_tags=samples[idx]["gt_tags"],
+                        error=f"Worker error: {e}",
+                    )
+                    results_by_index[idx] = result
+
+        # Reassemble in original order
+        results = [results_by_index[i] for i in range(total)]
 
     return results
 
@@ -506,6 +598,8 @@ def main(argv=None) -> int:
                     help="Use samples in file order (first N)")
     ap.add_argument("--seed", type=int, default=42,
                     help="Random seed for shuffle (default: 42)")
+    ap.add_argument("--workers", "-w", type=int, default=4,
+                    help="Number of parallel workers (default: 4, use 1 for sequential)")
 
     args = ap.parse_args(list(argv) if argv is not None else None)
 
@@ -522,6 +616,7 @@ def main(argv=None) -> int:
         verbose=args.verbose,
         shuffle=args.shuffle,
         seed=args.seed,
+        workers=args.workers,
     )
 
     print_summary(results)
@@ -551,6 +646,7 @@ def main(argv=None) -> int:
         "temperature": args.temperature,
         "shuffle": args.shuffle,
         "seed": args.seed,
+        "workers": args.workers,
         "n_errors": sum(1 for r in results if r.error),
     }