1"""Load raw per-attempt records into a tidy DataFrame with derived columns (scenario costs, method labels)."""2from __future__ import annotations3import glob, json4import pandas as pd5import numpy as np6from ..util import ROOT, read_jsonl78TEXT_METHODS = ["A0", "A1", "A2", "A3", "A4", "A5", "A6", "A7", "A8"]9CACHE_METHODS = ["B0", "B1", "B2", "B2b", "B3", "B4", "B5", "B6", "B9", "C0", "C2", "C3"]10LABELS = {"A0": "A0 no message", "A1": "A1 full text", "A2": "A2 head/tail", "A3": "A3 extractive", "A4": "A4 summary", "A5": "A5 structured evidence", "A6": "A6 refs+retrieval",11 "A7": "A7 zstd text", "A8": "A8 LLMLingua-2", "B0": "B0 native cache", "B1": "B1 serialized reuse", "B2": "B2 packed quant", "B3": "B3 sink+recent", "B4": "B4 random",12 "B5": "B5 attention (SnapKV-style)", "B6": "B6 pyramid", "B9": "B9 select+quant", "C0": "C0 identity (neg. control)", "C2": "C2 linear mapper", "C3": "C3 mapper+select"}1314def load_experiments(patterns: list[str]) -> pd.DataFrame:15 rows = []16 for pat in patterns:17 for p in glob.glob(str(ROOT / "runs/raw" / pat)):18 for r in read_jsonl(p):19 rows.append(flatten(r))20 df = pd.DataFrame(rows)21 return df2223def _scenario_costs(r: dict) -> dict:24 sp = r.get("timing_spans") or {}25 m = r["method_id"]26 enc = sp.get("encode_ms", 0.0) or 0.027 recv = (sp.get("receiver_prefill_ms", 0.0) or 0.0) + (sp.get("ttft_ms", 0.0) or 0.0) + (sp.get("decode_ms", 0.0) or 0.0)28 if m.startswith("B") or m.startswith("C"):29 doc_prefill = sp.get("doc_prefill_ms", 0.0) or 0.030 compress = (sp.get("compress_ms", 0.0) or 0.0) + (sp.get("serialize_ms", 0.0) or 0.0)31 load = sp.get("load_ms", 0.0) or 0.032 # cold: document must be read & cached, compressed, loaded, then query processed33 cold = doc_prefill + compress + load + recv34 # warm (cache already built & compressed once, reused across N queries): per-query = load + query35 warm_query = load + recv36 reuse_n = max(1, r.get("reuse_count") or 1)37 amortized = (doc_prefill + compress) / reuse_n + warm_query38 return {"cost_cold_ms": cold, "cost_warm_query_ms": warm_query, "cost_amortized_ms": amortized, "cost_build_ms": doc_prefill + compress, "receiver_ms": recv}39 else:40 reuse_n = max(1, r.get("reuse_count") or 1)41 cold = enc + recv42 warm_query = recv # message already built43 amortized = enc / reuse_n + recv44 return {"cost_cold_ms": cold, "cost_warm_query_ms": warm_query, "cost_amortized_ms": amortized, "cost_build_ms": enc, "receiver_ms": recv}4546def flatten(r: dict) -> dict:47 sc = r.get("task_scores") or {}48 cc = r.get("compression_config") or {}49 out = {k: r.get(k) for k in ["experiment_id", "trial_id", "attempt_id", "run_group_id", "task_family", "document_id", "question_id", "split", "cluster_id", "method_id",50 "question_access", "topology", "handoff_index", "reuse_count", "status", "error_code", "serialized_payload_bytes", "message_utf8_bytes",51 "cache_tensor_bytes", "fetched_bytes", "total_wall_ms", "ttft_ms", "peak_prefill_memory", "peak_decode_memory", "prediction_artifact", "protocol_hash"]}52 out["sender"] = (r.get("sender_model_revision") or "").split("@")[0].split("/")[-1]53 out["receiver"] = (r.get("receiver_model_revision") or "").split("@")[0].split("/")[-1]54 out["pair"] = f"{out['sender']}->{out['receiver']}" if out["sender"] else out["receiver"]55 out["correct"] = sc.get("correct", np.nan); out["f1"] = sc.get("f1", np.nan); out["em"] = sc.get("em", np.nan); out["pass_rate"] = sc.get("pass_rate", np.nan); out["abstained"] = sc.get("abstained", np.nan)56 out["quality"] = sc.get("f1") if r["task_family"].startswith("public_") else sc.get("correct")57 if r["status"] != "success":58 out["quality_operational"] = 0.059 else:60 out["quality_operational"] = out["quality"]61 out["budget_tokens"] = cc.get("budget_tokens"); out["bits"] = cc.get("bits"); out["keep_frac"] = cc.get("keep_frac"); out["codec"] = cc.get("codec")62 out["condition"] = r["run_group_id"].split(":", 1)[1] if ":" in (r.get("run_group_id") or "") else r["method_id"]63 out["label"] = LABELS.get(r["method_id"], r["method_id"])64 out["qtype"] = r["question_id"].split("_", 1)[1] if "_" in r["question_id"] else r["task_family"]65 ev = r.get("evidence_artifact") or {}66 if isinstance(ev, dict) and "ledger" in ev:67 L = ev["ledger"]; out["retention_rate"] = L["retention_rate"]; out["unsupported"] = L["unsupported_numbers"] + L["unsupported_identifiers"]; out["distorted"] = L["distorted_sentences"]; out["msg_bytes_hop"] = L["message_bytes"]68 out.update(_scenario_costs(r))69 sp = r.get("timing_spans") or {}70 for k in ["encode_ms", "doc_prefill_ms", "compress_ms", "serialize_ms", "load_ms", "receiver_prefill_ms", "query_prefill_ms", "decode_ms", "sender_prefill_ms"]:71 out["span_" + k] = sp.get(k)72 cost = r.get("cost_with_source") or {}73 out["cost_usd"] = cost.get("usd") if isinstance(cost, dict) else None74 return out75