PerturbReason/PerturbReason_dataset_code

1#!/usr/bin/env python32"""3eval_v3/export_rescue.py4=========================5Step 2 of the eval pipeline: export hard cases for LLM rescue.6 7Reads the Tier 1+2 samples CSV, filters samples with ambiguous error labels,8builds rescue prompts, and writes a JSONL file ready for batch vLLM inference.9 10Usage::11 12    ~/anaconda3/bin/python -m eval_v3.export_rescue \13        --samples-csv eval_v3/eval_v3_output/samples_*.csv \14        --pred-dir /path/to/noisy_context_output_0322 \15        --gt-dir /path/to/noisy_input \16        -o eval_v3/eval_v3_output/rescue_prompts.jsonl17"""18 19from __future__ import annotations20 21import argparse22import csv23import json24import sys25from pathlib import Path26from typing import Dict, List, Optional, Tuple27 28from .metric_llm import build_rescue_prompt, select_rescue_type29 30 31# Hard cases that need LLM rescue32HARD_LABELS = {33    "CORRECT_HALLUCINATED_EDGE",34    "CORRECT_RIGHT_FOR_WRONG_REASON",35    "WRONG_CORRECT_CHAIN",36}37 38 39def load_samples_csv(csv_path: Path) -> List[Dict[str, str]]:40    """Load the per-sample CSV produced by Tier 1+2 run."""41    rows = []42    with open(csv_path, "r", newline="") as f:43        reader = csv.DictReader(f)44        for row in reader:45            rows.append(row)46    return rows47 48 49def load_model_output(pred_dir: Path, split: str, file_name: str, idx: int) -> str:50    """Load the raw model_output for a given sample from the prediction JSONL."""51    pred_path = pred_dir / split / file_name52    if not pred_path.exists():53        return ""54    with open(pred_path, "r") as f:55        for i, line in enumerate(f):56            if i == idx:57                record = json.loads(line)58                return record.get("model_output", "")59    return ""60 61 62def build_file_index(pred_dir: Path) -> Dict[Tuple[str, str], Path]:63    """Build mapping from (split, file_name) → file path."""64    index = {}65    if not pred_dir.exists():66        return index67    for split_dir in sorted(pred_dir.iterdir()):68        if not split_dir.is_dir():69            continue70        for f in sorted(split_dir.glob("*.jsonl")):71            index[(split_dir.name, f.name)] = f72    return index73 74 75def load_pred_file_lines(pred_path: Path) -> List[dict]:76    """Load all lines from a prediction JSONL into a list."""77    records = []78    with open(pred_path, "r") as f:79        for line in f:80            line = line.strip()81            if line:82                records.append(json.loads(line))83    return records84 85 86def export_rescue_prompts(87    samples_csv: Path,88    pred_dir: Path,89    output_path: Path,90    max_samples: Optional[int] = None,91) -> int:92    """93    Read samples CSV, filter hard cases, build rescue prompts,94    write to output JSONL.95 96    Returns the number of rescue prompts written.97    """98    rows = load_samples_csv(samples_csv)99 100    # Group rows by (split, file_name) for efficient file reading101    from collections import defaultdict102    groups: Dict[Tuple[str, str], List[Tuple[int, Dict]]] = defaultdict(list)103    for i, row in enumerate(rows):104        error_label = row.get("error_label", "")105        if error_label in HARD_LABELS:106            key = (row["split"], row["file_name"])107            groups[key].append((i, row))108 109    # Also handle parse failures110    for i, row in enumerate(rows):111        if row.get("answer_parse_fail", "").lower() == "true":112            error_label = row.get("error_label", "")113            if error_label not in HARD_LABELS:  # don't double-count114                key = (row["split"], row["file_name"])115                groups[key].append((i, row))116 117    file_index = build_file_index(pred_dir)118 119    rescue_entries = []120    for (split, fname), items in sorted(groups.items()):121        pred_path = file_index.get((split, fname))122        if pred_path is None:123            print(f"  WARNING: prediction file not found for {split}/{fname}, skipping")124            continue125 126        pred_records = load_pred_file_lines(pred_path)127 128        for csv_row_idx, row in items:129            sample_id = row.get("sample_id", "")130            # sample_id is the positional index within the file131            try:132                pos_idx = int(sample_id)133            except (ValueError, TypeError):134                continue135 136            if pos_idx >= len(pred_records):137                continue138 139            model_output = pred_records[pos_idx].get("model_output", "")140 141            # Determine rescue type142            rescue_type = select_rescue_type(143                error_label=row.get("error_label", ""),144                answer_parse_fail=row.get("answer_parse_fail", "").lower() == "true",145                bleurt_score=_safe_float(row.get("bleurt_score")),146                answer_correct=row.get("answer_correct", "").lower() == "true",147            )148            if not rescue_type:149                continue150 151            # Build prompt152            prompt = build_rescue_prompt(153                rescue_type=rescue_type,154                model_output=model_output,155                perturbation=row.get("perturbation", ""),156                cell_type=row.get("cell_type", ""),157                effect_gene=row.get("effect_gene", ""),158                gt_answer=row.get("gt_answer", ""),159                model_answer=row.get("model_answer", ""),160                bleurt_score=_safe_float(row.get("bleurt_score")),161            )162 163            entry = {164                "prompt": prompt,165                "metadata": {166                    "csv_row_idx": csv_row_idx,167                    "sample_id": sample_id,168                    "split": split,169                    "file_name": fname,170                    "rescue_type": rescue_type,171                    "error_label": row.get("error_label", ""),172                    "gt_answer": row.get("gt_answer", ""),173                    "model_answer": row.get("model_answer", ""),174                    "perturbation": row.get("perturbation", ""),175                    "effect_gene": row.get("effect_gene", ""),176                    "cell_type": row.get("cell_type", ""),177                },178            }179            rescue_entries.append(entry)180 181    if max_samples is not None:182        rescue_entries = rescue_entries[:max_samples]183 184    # Write JSONL185    output_path.parent.mkdir(parents=True, exist_ok=True)186    with open(output_path, "w") as f:187        for entry in rescue_entries:188            f.write(json.dumps(entry, ensure_ascii=False) + "\n")189 190    return len(rescue_entries)191 192 193def _safe_float(val) -> Optional[float]:194    if val is None or val == "" or val == "None":195        return None196    try:197        return float(val)198    except (ValueError, TypeError):199        return None200 201 202def main():203    parser = argparse.ArgumentParser(204        description="Export hard cases for LLM rescue (Tier 4)")205    parser.add_argument("--samples-csv", type=Path, required=True,206                        help="Samples CSV from Tier 1+2 run")207    parser.add_argument("--pred-dir", type=Path, required=True,208                        help="Prediction directory (to read raw model outputs)")209    import time as _time210    _ts = _time.strftime("%Y%m%d_%H%M%S")211    parser.add_argument("-o", "--output", type=Path,212                        default=Path(__file__).parent / f"eval_v3_output/rescue_prompts_{_ts}.jsonl",213                        help="Output JSONL with rescue prompts")214    parser.add_argument("--max-samples", type=int, default=None,215                        help="Max rescue samples to export")216    args = parser.parse_args()217 218    print(f"Reading samples from {args.samples_csv}")219    n = export_rescue_prompts(220        samples_csv=args.samples_csv,221        pred_dir=args.pred_dir,222        output_path=args.output,223        max_samples=args.max_samples,224    )225    print(f"Exported {n} rescue prompts → {args.output}")226 227 228if __name__ == "__main__":229    main()230