PerturbReason/PerturbReason_dataset_code
012
1"""2eval_v3/pipeline.py3====================4Orchestrator — runs all evaluation tiers and produces unified output.5 6Handles the split directory structure:7 gt_dir/ → 5 splits × 2 pert types8 pred_dir/ → matching structure with qwen_pred_ prefix9 10Usage::11 12 from eval_v3.pipeline import EvalPipeline13 14 pipe = EvalPipeline(tiers=[1, 2])15 results = pipe.evaluate_all(gt_dir="noisy_input", pred_dir="noisy_output")16"""17 18from __future__ import annotations19 20import csv21import json22import time23from collections import defaultdict24from dataclasses import asdict25from pathlib import Path26from typing import Any, Dict, List, Optional, Set, Tuple27 28from .data_model import (29 FileResult, SampleRecord, SampleResult,30 load_paired_samples, load_external_kg, find_file_pairs,31)32from .metric_accuracy import evaluate_accuracy33from .metric_symbolic import evaluate_symbolic34from .metric_go_similarity import GoAnnotations, evaluate_go_similarity35 36 37class EvalPipeline:38 """39 Configurable multi-tier evaluation pipeline.40 41 Parameters42 ----------43 tiers : which metric tiers to run (1-3).44 go_gmt_path : path to MSigDB GMT file for Tier 3 GO similarity.45 external_kg_path : path to external INDRA KG JSON.46 """47 48 def __init__(49 self,50 tiers: List[int] | None = None,51 go_gmt_path: str | Path | None = None,52 external_kg_path: str | Path | None = None,53 ):54 self.tiers = set(tiers or [1, 2])55 56 # Load GO annotations once (Tier 3)57 self.go_annotations: Optional[GoAnnotations] = None58 if go_gmt_path and Path(go_gmt_path).exists():59 print(f"[Pipeline] Loading GO GMT from {go_gmt_path} ...")60 anno = GoAnnotations()61 anno.load_gmt(go_gmt_path)62 anno.finalize()63 self.go_annotations = anno64 print(f"[Pipeline] GO annotations: {len(anno.gene_to_terms)} genes, "65 f"{len(anno.idf)} terms")66 elif 3 in set(tiers or [1, 2]):67 print("[Pipeline] Warning: --go-gmt not provided; Tier 3 will be skipped.")68 69 # Load external KG once70 self.external_kg_edges: Set[Tuple[str, str, str]] = set()71 self.external_kg_pairs: Set[Tuple[str, str]] = set()72 if external_kg_path and Path(external_kg_path).exists():73 print(f"[Pipeline] Loading external KG from {external_kg_path} ...")74 self.external_kg_edges, self.external_kg_pairs = load_external_kg(75 external_kg_path)76 print(f"[Pipeline] External KG: {len(self.external_kg_edges)} edges, "77 f"{len(self.external_kg_pairs)} unique pairs")78 79 # ──────────────────────────────────────────80 # Single file-pair evaluation81 # ──────────────────────────────────────────82 83 def evaluate_file_pair(84 self,85 gt_path: str | Path,86 pred_path: str | Path,87 split_name: str = "",88 pert_type: str = "",89 ) -> FileResult:90 """91 Evaluate one GT/prediction file pair through all configured tiers.92 """93 gt_path = Path(gt_path)94 pred_path = Path(pred_path)95 96 samples = load_paired_samples(gt_path, pred_path)97 98 # Initialize results99 results = [100 SampleResult(101 id=s.id,102 cell_type=s.prompt_data.cell_type if s.prompt_data else "",103 pert_type=s.prompt_data.pert_type if s.prompt_data else s.pert_type,104 perturbation=s.prompt_data.perturbation if s.prompt_data else "",105 effect_gene=s.prompt_data.effect_gene if s.prompt_data else "",106 task="3way",107 )108 for s in samples109 ]110 111 # ── Tier 1: Accuracy ──112 tier1_scores: Dict[str, Any] = {}113 if 1 in self.tiers:114 tier1_scores = evaluate_accuracy(samples, results)115 116 # ── Tier 2: Symbolic ──117 tier2_scores: Dict[str, Any] = {}118 if 2 in self.tiers:119 tier2_scores = evaluate_symbolic(120 samples, results,121 external_kg_edges=self.external_kg_edges,122 external_kg_pairs=self.external_kg_pairs,123 )124 125 # ── Tier 3: GO Similarity ──126 tier3_scores: Dict[str, Any] = {}127 if 3 in self.tiers:128 tier3_scores = evaluate_go_similarity(129 samples, results,130 go_annotations=self.go_annotations,131 )132 133 # Tier 4 (LLM Rescue) is run separately via:134 # export_rescue.py → qwen_rescue_vllm.py → merge_rescue.py135 136 aggregate = {137 "tier1_accuracy": tier1_scores,138 "tier2_symbolic": tier2_scores,139 "tier3_go_sim": tier3_scores,140 }141 142 return FileResult(143 file_name=pred_path.name,144 split_name=split_name,145 pert_type=pert_type,146 metadata={147 "split": split_name,148 "pert_type": pert_type,149 "gt_file": gt_path.name,150 "pred_file": pred_path.name,151 },152 sample_results=results,153 aggregate=aggregate,154 )155 156 # ──────────────────────────────────────────157 # Full directory evaluation158 # ──────────────────────────────────────────159 160 def evaluate_all(161 self,162 gt_dir: str | Path,163 pred_dir: str | Path,164 ) -> List[FileResult]:165 """166 Evaluate all matching file pairs across the split directory structure.167 """168 pairs = find_file_pairs(gt_dir, pred_dir)169 file_results: List[FileResult] = []170 171 print(f"\n[Pipeline] Found {len(pairs)} file pairs to evaluate")172 print(f"[Pipeline] Tiers: {sorted(self.tiers)}")173 print()174 175 for gt_path, pred_path, split_name, pert_type in pairs:176 print(f" Evaluating {split_name}/{pert_type}: {pred_path.name} ...")177 try:178 fr = self.evaluate_file_pair(179 gt_path, pred_path, split_name, pert_type)180 file_results.append(fr)181 182 t1 = fr.aggregate.get("tier1_accuracy", {})183 acc = t1.get("accuracy", "N/A")184 bal = t1.get("balanced_accuracy", "N/A")185 n = len(fr.sample_results)186 if isinstance(acc, float):187 print(f" → {n} samples, accuracy={acc:.4f}, "188 f"balanced_acc={bal:.4f}")189 else:190 print(f" → {n} samples")191 except Exception as e:192 print(f" ERROR: {e}")193 import traceback194 traceback.print_exc()195 196 return file_results197 198 # ──────────────────────────────────────────199 # Output writers200 # ──────────────────────────────────────────201 202 @staticmethod203 def write_summary_json(204 file_results: List[FileResult], output_path: str | Path,205 ):206 """Write aggregate scores for all files to a JSON."""207 output_path = Path(output_path)208 output_path.parent.mkdir(parents=True, exist_ok=True)209 data = {}210 for fr in file_results:211 key = f"{fr.split_name}/{fr.pert_type}"212 data[key] = {213 "file_name": fr.file_name,214 "metadata": fr.metadata,215 "aggregate": fr.aggregate,216 "num_samples": len(fr.sample_results),217 }218 with open(output_path, "w") as f:219 json.dump(data, f, indent=2, default=str)220 221 @staticmethod222 def write_samples_csv(223 file_results: List[FileResult], output_path: str | Path,224 ):225 """Write per-sample results across all files to a CSV."""226 output_path = Path(output_path)227 output_path.parent.mkdir(parents=True, exist_ok=True)228 229 all_rows: List[Dict[str, Any]] = []230 for fr in file_results:231 for sr in fr.sample_results:232 row = {233 "split": fr.split_name,234 "file_pert_type": fr.pert_type,235 "file_name": fr.file_name,236 "sample_id": sr.id,237 "cell_type": sr.cell_type,238 "pert_type": sr.pert_type,239 "perturbation": sr.perturbation,240 "effect_gene": sr.effect_gene,241 "task": sr.task,242 # Tier 1243 "gt_answer": sr.gt_answer,244 "model_answer": sr.model_answer,245 "answer_correct": sr.answer_correct,246 "answer_parse_fail": sr.answer_parse_fail,247 # Tier 2248 "edge_f1_strict": sr.edge_f1_strict,249 "edge_f1_relaxed": sr.edge_f1_relaxed,250 "edge_recall_strict": sr.edge_recall_strict,251 "edge_precision_strict": sr.edge_precision_strict,252 "path_connectivity": sr.path_connectivity,253 "sign_match_rate": sr.sign_match_rate,254 "sign_flips": sr.sign_flips,255 "in_kg_rate": sr.in_kg_rate,256 "hallucination_rate": sr.hallucination_rate,257 "num_model_triplets": sr.num_model_triplets,258 "num_hallucinated_edges": sr.num_hallucinated_edges,259 "error_label": sr.error_label,260 # Tier 3261 "go_sim_score": sr.go_sim_score,262 # Tier 4263 "llm_rescue_label": sr.llm_rescue_label,264 "llm_rescue_applied": sr.llm_rescue_applied,265 }266 all_rows.append(row)267 268 if not all_rows:269 return270 271 with open(output_path, "w", newline="") as f:272 writer = csv.DictWriter(f, fieldnames=all_rows[0].keys())273 writer.writeheader()274 writer.writerows(all_rows)275 276 @staticmethod277 def write_file_summary_csv(278 file_results: List[FileResult], output_path: str | Path,279 ):280 """Write one row per file with aggregate scores."""281 output_path = Path(output_path)282 output_path.parent.mkdir(parents=True, exist_ok=True)283 284 rows: List[Dict[str, Any]] = []285 for fr in file_results:286 t1 = fr.aggregate.get("tier1_accuracy", {})287 t2 = fr.aggregate.get("tier2_symbolic", {})288 t3 = fr.aggregate.get("tier3_go_sim", {})289 t4 = fr.aggregate.get("tier4_llm_rescue", {})290 row: Dict[str, Any] = {291 "split": fr.split_name,292 "pert_type": fr.pert_type,293 "file_name": fr.file_name,294 "num_samples": len(fr.sample_results),295 # T1296 "accuracy": t1.get("accuracy"),297 "balanced_accuracy": t1.get("balanced_accuracy"),298 "f1_macro": t1.get("f1_macro"),299 "f1_weighted": t1.get("f1_weighted"),300 "parse_failures": t1.get("parse_failures"),301 # T2302 "edge_f1_strict_mean": t2.get("edge_f1_strict_mean"),303 "edge_f1_relaxed_mean": t2.get("edge_f1_relaxed_mean"),304 "edge_recall_strict_mean": t2.get("edge_recall_strict_mean"),305 "edge_precision_strict_mean": t2.get("edge_precision_strict_mean"),306 "path_connected_rate": t2.get("path_connected_rate"),307 "sign_match_rate_mean": t2.get("sign_match_rate_mean"),308 "in_kg_rate_mean": t2.get("in_kg_rate_mean"),309 "hallucination_rate_mean": t2.get("hallucination_rate_mean"),310 # T3311 "go_sim_mean": t3.get("go_sim_mean"),312 "go_sim_median": t3.get("go_sim_median"),313 # T4314 "llm_rescue_total_unsolved": t4.get("total_unsolved"),315 "llm_rescue_total_processed": t4.get("total_processed"),316 }317 # Error taxonomy counts318 for et, count in t2.get("error_label_counts", {}).items():319 row[f"err_{et}"] = count320 # Per-class accuracy321 for lab, vals in t1.get("per_class", {}).items():322 row[f"precision_{lab}"] = vals["precision"]323 row[f"recall_{lab}"] = vals["recall"]324 row[f"f1_{lab}"] = vals["f1"]325 rows.append(row)326 327 if not rows:328 return329 330 all_cols: List[str] = list(rows[0].keys())331 for r in rows:332 for k in r:333 if k not in all_cols:334 all_cols.append(k)335 336 with open(output_path, "w", newline="") as f:337 writer = csv.DictWriter(f, fieldnames=all_cols, extrasaction="ignore")338 writer.writeheader()339 writer.writerows(rows)340 341 @staticmethod342 def print_report(fr: FileResult) -> Dict[str, Any]:343 """Pretty-print evaluation report for one file."""344 t1 = fr.aggregate.get("tier1_accuracy", {})345 t2 = fr.aggregate.get("tier2_symbolic", {})346 t3 = fr.aggregate.get("tier3_go_sim", {})347 t4 = fr.aggregate.get("tier4_llm_rescue", {})348 349 n = len(fr.sample_results)350 print(f"\n{'=' * 72}")351 print(f" Split: {fr.split_name} | Pert type: {fr.pert_type}")352 print(f" File: {fr.file_name}")353 print(f"{'=' * 72}")354 355 # Tier 1356 if t1:357 acc = t1.get("accuracy", 0)358 bal_acc = t1.get("balanced_accuracy", 0)359 f1_m = t1.get("f1_macro", 0)360 pf = t1.get("parse_failures", 0)361 print(f"\n [Tier-1] Answer Accuracy")362 print(f" Samples: {n}")363 print(f" Accuracy: {acc:.4f}")364 print(f" Balanced Accuracy: {bal_acc:.4f}")365 print(f" F1 (macro): {f1_m:.4f}")366 print(f" Parse failures: {pf}")367 pc = t1.get("per_class", {})368 if pc:369 for lab, vals in sorted(pc.items()):370 print(f" {lab}: P={vals['precision']:.3f} "371 f"R={vals['recall']:.3f} F1={vals['f1']:.3f}")372 373 # Tier 2374 if t2:375 print(f"\n [Tier-2] Symbolic Reasoning")376 print(f" Edge F1 (strict): {t2.get('edge_f1_strict_mean', 0):.4f}")377 print(f" Edge F1 (relaxed): {t2.get('edge_f1_relaxed_mean', 0):.4f}")378 print(f" Edge Recall: {t2.get('edge_recall_strict_mean', 0):.4f}")379 print(f" Edge Precision: {t2.get('edge_precision_strict_mean', 0):.4f}")380 print(f" Path Connected: {t2.get('path_connected_rate', 0):.4f}")381 print(f" Sign Match Rate: {t2.get('sign_match_rate_mean', 0):.4f}")382 print(f" In-KG Rate: {t2.get('in_kg_rate_mean', 0):.4f}")383 print(f" Hallucination Rate: {t2.get('hallucination_rate_mean', 0):.4f}")384 print(f" External KG loaded: {t2.get('external_kg_loaded', False)}")385 etc = t2.get("error_label_counts", {})386 if etc:387 print(f" Error distribution:")388 for et, cnt in sorted(etc.items()):389 rate = t2.get("error_label_rates", {}).get(et, 0)390 print(f" {et:<40s} {cnt:>5d} ({rate:.1%})")391 392 # Tier 3393 if t3 and t3.get("go_sim_available"):394 print(f"\n [Tier-3] GO Functional Similarity")395 print(f" Mean: {t3.get('go_sim_mean', 'N/A')}")396 print(f" Median: {t3.get('go_sim_median', 'N/A')}")397 print(f" Scored: {t3.get('num_scored', 'N/A')}")398 print(f" No GO cover: {t3.get('num_no_go_coverage', 'N/A')}")399 400 # Tier 4401 if t4 and t4.get("llm_rescue_available"):402 print(f"\n [Tier-4] LLM Rescue")403 print(f" Unsolved cases: {t4.get('total_unsolved', 0)}")404 print(f" Processed: {t4.get('total_processed', 0)}")405 rtypes = t4.get("rescue_type_counts", {})406 for rt, cnt in rtypes.items():407 print(f" {rt:<30s} {cnt:>4d}")408 409 print()410 return fr.aggregate411