PerturbReason/PerturbReason_dataset_code

1"""2eval_v3/data_model.py3=====================4Shared data structures, IO utilities, and parsing for the PerturbQA v3 pipeline.5 6Key changes from eval_v2:7  - Model outputs now contain structured <thinking>, <answer>, <triplet> blocks8  - GT label is in the input file's "label" field9  - GT triplets are in input file's "output.triplets_cell_conditioned"10  - No external LLM-based triplet extraction needed11"""12 13from __future__ import annotations14 15import json16import re17from collections import defaultdict18from dataclasses import dataclass, field, asdict19from pathlib import Path20from typing import Any, Dict, List, Optional, Set, Tuple21 22 23# ────────────────────────────────────────────24# Causal-edge & causal-path25# ────────────────────────────────────────────26 27@dataclass(frozen=True)28class CausalEdge:29    """A single directed edge in the biological causal graph."""30    source: str31    target: str32    sign: str  # "activates", "inhibits", "regulates"33 34    @property35    def numeric(self) -> Optional[int]:36        return {"activates": +1, "inhibits": -1}.get(self.sign)37 38    def __repr__(self) -> str:39        sym = {"activates": "─(+)→", "inhibits": "─(-)→", "regulates": "─(?)→"}40        return f"{self.source} {sym.get(self.sign, '→')} {self.target}"41 42 43# ────────────────────────────────────────────44# Prompt metadata (parsed from prompt text)45# ────────────────────────────────────────────46 47@dataclass48class PromptData:49    """Structured content parsed from the prompt text."""50    cell_type: str51    pert_type: str          # "chemical" | "genetic"52    perturbation: str       # drug name or perturbation description53    pert_gene: str          # actual gene name for genetic perturbation54    effect_gene: str55    task: str               # always "3way"56    basal_expression: Dict[str, str] = field(default_factory=dict)57    edges: List[CausalEdge] = field(default_factory=list)58    all_genes: Set[str] = field(default_factory=set)59 60 61# ────────────────────────────────────────────62# Sample record (one prediction instance)63# ────────────────────────────────────────────64 65@dataclass66class SampleRecord:67    """68    One evaluation sample, combining GT data (from input file)69    and model prediction (from output file).70 71    Fields72    ------73    id : sample identifier (positional index in JSONL)74    prompt : full prompt text75    model_output : raw model generation (contains <thinking>, <answer>, <triplet>)76    gt_label : ground-truth answer ("up", "down", "unchanged")77    gt_triplets : GT causal path from triplets_cell_conditioned78    gt_reasoning : GT reasoning text (if available)79    prompt_data : structured parsed metadata from prompt80    basal_context : per-gene basal expression values81    pert_type : "chemical" or "genetic"82    """83    id: Any84    prompt: str85    model_output: str86    gt_label: str87    gt_triplets: List[Tuple[str, str, str]] = field(default_factory=list)88    gt_reasoning: str = ""89    prompt_data: Optional[PromptData] = None90    basal_context: Dict[str, Any] = field(default_factory=dict)91    pert_type: str = ""92 93 94# ────────────────────────────────────────────95# Per-sample result container96# ────────────────────────────────────────────97 98@dataclass99class SampleResult:100    """Holds the evaluation result for a single sample across all tiers."""101    id: Any102    # ── Metadata ──103    cell_type: str = ""104    pert_type: str = ""105    perturbation: str = ""106    effect_gene: str = ""107    task: str = ""108 109    # ── Tier-1: Answer accuracy ──110    gt_answer: str = ""111    model_answer: str = ""112    answer_correct: Optional[bool] = None113    answer_parse_fail: bool = False114 115    # ── Tier-2: Symbolic reasoning ──116    # M1: Edge F1117    edge_f1_strict: float = 0.0118    edge_f1_relaxed: float = 0.0119    edge_recall_strict: float = 0.0120    edge_precision_strict: float = 0.0121    # M2: Path Connectivity122    path_connectivity: str = ""123    # M3: Sign Consistency124    sign_match_rate: float = 0.0125    sign_flips: int = 0126    # M4: Graph Validity127    in_kg_rate: float = 0.0128    hallucination_rate: float = 0.0129    num_model_triplets: int = 0130    num_hallucinated_edges: int = 0131    # M5: Error Taxonomy132    error_label: str = ""133 134    # ── Tier-3: GO Functional Similarity ──135    go_sim_score: Optional[float] = None136 137    # ── Tier-4: LLM Rescue ──138    llm_rescue_label: str = ""139    llm_rescue_applied: bool = False140    llm_rescue_explanation: str = ""141 142    def to_dict(self) -> Dict[str, Any]:143        return asdict(self)144 145 146# ────────────────────────────────────────────147# File-level aggregate148# ────────────────────────────────────────────149 150@dataclass151class FileResult:152    """Aggregate of all SampleResult values for one prediction file."""153    file_name: str154    split_name: str = ""155    pert_type: str = ""156    metadata: Dict[str, str] = field(default_factory=dict)157    sample_results: List[SampleResult] = field(default_factory=list)158    aggregate: Dict[str, Any] = field(default_factory=dict)159 160 161# ────────────────────────────────────────────162# Structured output parsing163# ────────────────────────────────────────────164 165def extract_thinking(text: str) -> str:166    """Extract text from <thinking>...</thinking> block."""167    m = re.search(r"<thinking>(.*?)</thinking>", text, re.DOTALL)168    return m.group(1).strip() if m else ""169 170 171def extract_answer(text: str) -> Optional[str]:172    """Extract and normalize the answer from <answer>...</answer> block.173 174    When models self-correct, they may emit multiple <answer> blocks.175    We use the **last** valid <answer> block, which is the corrected one.176    """177    valid = {"up", "down", "unchanged"}178 179    # Primary: LAST <answer>...</answer>180    matches = re.findall(r"<answer>\s*(\w+)\s*</answer>", text, re.IGNORECASE)181    for ans_raw in reversed(matches):182        ans = ans_raw.strip().lower()183        if ans in valid:184            return ans185 186    # Fallback: \boxed{...}187    m = re.search(r"\\boxed\{(\w+)\}", text)188    if m:189        ans = m.group(1).strip().lower()190        if ans in valid:191            return ans192 193    # Fallback: last word194    text_clean = text.strip()195    if text_clean:196        last_word = text_clean.rstrip(".").split()[-1].lower()197        if last_word in valid:198            return last_word199 200    return None201 202 203def extract_triplets(text: str) -> List[Tuple[str, str, str]]:204    """205    Extract causal triplets from <triplet>...</triplet> block.206 207    When models self-correct, they may emit multiple <triplet> blocks.208    We use the **last** block, which is the corrected one.209 210    Handles both proper JSON arrays and common formatting variations211    (missing outer brackets, trailing commas, etc.).212    """213    blocks = re.findall(r"<triplet>(.*?)</triplet>", text, re.DOTALL)214    if not blocks:215        return []216 217    # Try each block from last to first until one parses successfully218    for raw in reversed(blocks):219        raw = raw.strip()220        if not raw:221            continue222 223        result = _try_parse_triplet_block(raw)224        if result:225            return result226 227    return []228 229 230def _try_parse_triplet_block(raw: str) -> List[Tuple[str, str, str]]:231    """Attempt to parse a single raw triplet block string."""232    # Try parsing as proper JSON array of arrays233    try:234        parsed = json.loads(raw)235        if isinstance(parsed, list):236            result = _validate_triplets(parsed)237            if result:238                return result239    except json.JSONDecodeError:240        pass241 242    # Common issue: list of lists without outer brackets243    # e.g. ["A", "inhibits", "B"],\n["C", "activates", "D"]244    if not raw.startswith("[["):245        wrapped = "[" + raw.rstrip(",").rstrip() + "]"246        try:247            parsed = json.loads(wrapped)248            if isinstance(parsed, list):249                result = _validate_triplets(parsed)250                if result:251                    return result252        except json.JSONDecodeError:253            pass254 255    # Regex fallback: extract individual triplets256    triplets = []257    for match in re.finditer(258        r'\[\s*"([^"]+)"\s*,\s*"([^"]+)"\s*,\s*"([^"]+)"\s*\]', raw259    ):260        triplets.append((match.group(1), _normalize_relation(match.group(2)), match.group(3)))261    return [t for t in triplets if t[1] is not None]262 263 264# ── Relation normalization map ──265# Maps non-standard relation strings to canonical {activates, inhibits, regulates}.266_RELATION_MAP = {267    # Standard (identity)268    "activates": "activates",269    "activate": "activates",270    "inhibits": "inhibits",271    "regulates": "regulates",272    "regulate": "regulates",273    # Positive / activating synonyms274    "upregulates": "activates",275    "increases": "activates",276    "promotes": "activates",277    "induces": "activates",278    "stabilizes": "activates",279    "supports": "activates",280    "positively regulates": "activates",281    "indirectly activates": "activates",282    "induces upregulation of": "activates",283    "phosphorylates": "activates",284    "is upregulated": "activates",285    # Negative / inhibiting synonyms286    "suppresses": "inhibits",287    "downregulates": "inhibits",288    "reduces": "inhibits",289    "represses": "inhibits",290    "inactivates": "inhibits",291    "negatively regulates": "inhibits",292    "indirectly suppresses": "inhibits",293    "indirectly reduces": "inhibits",294    "phosphorylates and destabilizes": "inhibits",295    "inhibited by": "inhibits",296    "inhibited_by": "inhibits",297    # Ambiguous → regulates298    "modulates": "regulates",299    "influences": "regulates",300    "affects": "regulates",301    "indirectly regulates": "regulates",302    "leads to": "regulates",303    "involved in": "regulates",304    "responds to": "regulates",305    "regulated by": "regulates",306    "regulated_by": "regulates",307    "is regulated by": "regulates",308    "upregulated by": "regulates",309    "downregulated by": "regulates",310    "repressed by": "regulates",311    "is stabilized by": "regulates",312    "is not regulated by": None,   # negation → drop313    "not regulated": None,         # negation → drop314    "does not regulate": None,     # negation → drop315    "no regulatory link": None,    # negation → drop316    "expression": "regulates",317}318 319 320def _normalize_relation(rel: str) -> Optional[str]:321    """Map a relation string to canonical form, or None if unrecognized."""322    return _RELATION_MAP.get(rel.lower().strip())323 324 325def _validate_triplets(parsed: list) -> List[Tuple[str, str, str]]:326    """Validate and convert parsed JSON to list of triplet tuples.327 328    Non-standard relations (e.g. 'suppresses', 'upregulates') are329    normalized to the canonical set {activates, inhibits, regulates}.330    Triplets with truly unknown relations are discarded.331    """332    result = []333    for item in parsed:334        if isinstance(item, (list, tuple)) and len(item) == 3:335            src, tgt = str(item[0]), str(item[2])336            norm_rel = _normalize_relation(str(item[1]))337            if norm_rel is not None:338                result.append((src, norm_rel, tgt))339    return result340 341 342# ────────────────────────────────────────────343# Prompt parsing344# ────────────────────────────────────────────345 346def parse_prompt(prompt: str) -> PromptData:347    """Extract structured info from a PerturbQA prompt string."""348    cell = re.search(r"\*\*Cell Type:\*\*\s*(.+)", prompt)349    pert = re.search(r"\*\*Perturbation \((\w+)\):\*\*\s*(.+)", prompt)350    gene = re.search(r"\*\*Effect Gene:\*\*\s*(.+)", prompt)351 352    # Basal expression353    basal: Dict[str, str] = {}354    basal_sec = re.search(355        r"basal gene expression.*?:\n(.*?)(?=\n##)", prompt, re.DOTALL | re.IGNORECASE356    )357    if basal_sec:358        for m in re.finditer(r"-\s+(\S+):\s*(\w+)", basal_sec.group(1)):359            basal[m.group(1)] = m.group(2)360 361    # Knowledge edges362    edges: List[CausalEdge] = []363    know_sec = re.search(r"## Retrieved Knowledges?\n(.*?)(?=\n##)", prompt, re.DOTALL)364    if know_sec:365        for line in know_sec.group(1).split("\n"):366            line = line.strip().lstrip("- ")367            if line:368                e = _parse_edge_text(line)369                if e:370                    edges.append(e)371 372    all_genes: Set[str] = set()373    for e in edges:374        all_genes.add(e.source)375        all_genes.add(e.target)376    all_genes.update(basal.keys())377 378    pert_name = pert.group(2).strip() if pert else ""379    pert_type_str = pert.group(1).strip() if pert else ""380    effect_name = gene.group(1).strip() if gene else ""381 382    # For genetic perturbations, extract gene name383    pert_gene = ""384    if pert_type_str.lower() == "genetic":385        gene_match = re.search(386            r"(?:knockdown|knockout|overexpression|activation)\s+(?:of\s+)?(\S+)",387            pert_name, re.IGNORECASE,388        )389        if gene_match:390            pert_gene = gene_match.group(1)391        else:392            parts = pert_name.split()393            if parts:394                pert_gene = parts[-1]395    else:396        pert_gene = pert_name397 398    if pert_name:399        all_genes.add(pert_name)400    if pert_gene and pert_gene != pert_name:401        all_genes.add(pert_gene)402    if effect_name:403        all_genes.add(effect_name)404 405    return PromptData(406        cell_type=cell.group(1).strip() if cell else "",407        pert_type=pert_type_str,408        perturbation=pert_name,409        pert_gene=pert_gene,410        effect_gene=effect_name,411        task="3way",412        basal_expression=basal,413        edges=edges,414        all_genes=all_genes,415    )416 417 418def _parse_edge_text(text: str) -> Optional[CausalEdge]:419    """Parse 'A inhibits B' → CausalEdge."""420    for pat, sign in [421        (r"(.+?)\s+inhibits?\s+(.+)", "inhibits"),422        (r"(.+?)\s+activates?\s+(.+)", "activates"),423        (r"(.+?)\s+regulates?\s+(.+)", "regulates"),424    ]:425        m = re.match(pat, text.strip(), re.IGNORECASE)426        if m:427            return CausalEdge(m.group(1).strip(), m.group(2).strip(), sign)428    return None429 430 431# ────────────────────────────────────────────432# External KG loading433# ────────────────────────────────────────────434 435def load_external_kg(436    kg_path: Path | str,437) -> Tuple[Set[Tuple[str, str, str]], Set[Tuple[str, str]]]:438    """439    Load an external KG from a PKL (NetworkX DiGraph) or JSON file.440 441    PKL format: networkx.DiGraph where each edge has attribute442        {'relation': 'activates' | 'inhibits' | ...}443 444    JSON format: {"edges": [[src, rel, tgt], ...]}445 446    Returns (edge_set, pair_set) where:447        edge_set = set of (src, rel, tgt) tuples448        pair_set = set of (src, tgt) tuples (for relaxed matching)449    """450    kg_path = Path(kg_path)451    edge_set: Set[Tuple[str, str, str]] = set()452    pair_set: Set[Tuple[str, str]] = set()453 454    if kg_path.suffix == ".pkl":455        import pickle456        with open(kg_path, "rb") as f:457            G = pickle.load(f)458        for src, tgt, attr in G.edges(data=True):459            rel = attr.get("relation", "regulates")460            edge_set.add((str(src), rel, str(tgt)))461            pair_set.add((str(src), str(tgt)))462    else:463        with open(kg_path, "r", encoding="utf-8") as f:464            data = json.load(f)465        for e in data.get("edges", []):466            if len(e) >= 3:467                src, rel, tgt = e[0], e[1], e[2]468                edge_set.add((src, rel, tgt))469                pair_set.add((src, tgt))470 471    return edge_set, pair_set472 473 474# ────────────────────────────────────────────475# IO helpers476# ────────────────────────────────────────────477 478def load_jsonl(path: Path | str) -> List[Dict[str, Any]]:479    """Load a JSONL file, returning a list of dicts."""480    records: List[Dict[str, Any]] = []481    with open(path, "r", encoding="utf-8") as f:482        for line in f:483            line = line.strip()484            if line:485                records.append(json.loads(line))486    return records487 488 489def load_paired_samples(490    gt_path: Path | str,491    pred_path: Path | str,492) -> List[SampleRecord]:493    """494    Load paired GT + prediction files into SampleRecords.495 496    GT file format (noisy_input):497      {id, label, pert_type, prompt, response, input, output, basal_context, data_type}498      - output.triplets_cell_conditioned = GT causal path499      - label = GT answer (up/down/unchanged)500 501    Prediction file format (noisy_context_output):502      {source_file, prompt, ground_truth_response, model_output}503      - model_output = <thinking>...<answer>...<triplet>...504    """505    gt_records = load_jsonl(gt_path)506    pred_records = load_jsonl(pred_path)507 508    # Match by position (both files have same ordering)509    if len(gt_records) != len(pred_records):510        print(f"  WARNING: GT has {len(gt_records)} samples, "511              f"predictions has {len(pred_records)} samples. "512              f"Using min({len(gt_records)}, {len(pred_records)}).")513 514    n = min(len(gt_records), len(pred_records))515    samples: List[SampleRecord] = []516 517    for i in range(n):518        gt = gt_records[i]519        pred = pred_records[i]520 521        prompt_text = gt.get("prompt", "") or pred.get("prompt", "")522        pd = parse_prompt(prompt_text)523 524        # GT data525        gt_label = gt.get("label", "")526        gt_output = gt.get("output", {}) or {}527        gt_triplets_raw = gt_output.get("triplets_cell_conditioned", [])528        gt_triplets = [tuple(t) for t in gt_triplets_raw if len(t) == 3]529        basal = gt.get("basal_context", {}) or {}530        pert_type = gt.get("pert_type", "")531 532        # Model output533        model_output = pred.get("model_output", "")534 535        rec = SampleRecord(536            id=gt.get("id", i),537            prompt=prompt_text,538            model_output=model_output,539            gt_label=gt_label,540            gt_triplets=gt_triplets,541            prompt_data=pd,542            basal_context=basal,543            pert_type=pert_type,544        )545        samples.append(rec)546 547    return samples548 549 550def _classify_pert_type(fname_lower: str) -> Optional[str]:551    """552    Classify a JSONL filename into a pert_type string.553 554    More-specific compound keywords are checked first to avoid the generic555    "chemical" / "genetic" matches swallowing them.556 557    Compound types (used in other_test):558      chemical2genetic, genetic2chemical, genetic_combo559 560    Generic types (all other splits):561      chemical, genetic562    """563    # Compound keywords — must be checked before the generic ones564    for keyword in ("chemical2genetic", "genetic2chemical", "genetic_combo"):565        if keyword in fname_lower:566            return keyword567    if "chemical" in fname_lower:568        return "chemical"569    if "genetic" in fname_lower:570        return "genetic"571    return None572 573 574def find_file_pairs(575    gt_dir: Path | str,576    pred_dir: Path | str,577) -> List[Tuple[Path, Path, str, str]]:578    """579    Find matching GT/prediction file pairs across the split directory structure.580 581    Both directories must share the same 5 split sub-folders.  Within each582    split folder JSONL files are classified by pert_type keyword:583 584    * Generic splits (id_test, cell_ood, pert_ood, double_ood):585        exactly one file with "chemical" and one with "genetic".586    * other_test (extended):587        up to 3 files — "chemical2genetic", "genetic2chemical", "genetic_combo"588        (and still supports the legacy 2-file layout for backwards compat).589 590    Files are paired by split + pert_type — no filename prefix conventions591    are assumed.592 593    Returns list of (gt_path, pred_path, split_name, pert_type).594    """595    gt_dir = Path(gt_dir)596    pred_dir = Path(pred_dir)597    pairs: List[Tuple[Path, Path, str, str]] = []598 599    for split_dir in sorted(gt_dir.iterdir()):600        if not split_dir.is_dir():601            continue602        split_name = split_dir.name603 604        pred_split_dir = pred_dir / split_name605        if not pred_split_dir.is_dir():606            print(f"  WARNING: No prediction directory for split '{split_name}'")607            continue608 609        # Index pred files by pert_type keyword610        pred_by_type: Dict[str, Path] = {}611        for f in pred_split_dir.glob("*.jsonl"):612            pt = _classify_pert_type(f.name.lower())613            if pt is not None:614                pred_by_type[pt] = f615 616        for gt_file in sorted(split_dir.glob("*.jsonl")):617            pert_type = _classify_pert_type(gt_file.name.lower())618            if pert_type is None:619                print(f"  WARNING: Cannot determine pert_type for {gt_file.name}, skipping")620                continue621 622            pred_file = pred_by_type.get(pert_type)623            if pred_file is None:624                print(f"  WARNING: No '{pert_type}' prediction file in {pred_split_dir}")625                continue626 627            pairs.append((gt_file, pred_file, split_name, pert_type))628 629    return pairs630