1"""E5: fan-in / composition with conflicting evidence.23The synthetic document is partitioned into K source agents (project blocks round-robin). Overlap: one project block is4duplicated into a second source with an *earlier* (superseded) revision sentence, so the ledger's precedence rule still5yields the original answers but a naive merge sees a conflict. One source additionally receives an irrelevant block from6another document. The aggregator answers multi-hop/aggregation/revision questions that need several sources.78Aggregation methods: concat_A4 (concatenated prose summaries), merge_A5 (merged structured records), shared_memory_A69(reference index per source + metered retrieval by the aggregator), full_context (reference: all source texts concatenated),10and cache_concat_negative_control (KV caches of independently prefilled sources concatenated along the position axis —11positions overlap and no cross-source context exists; reported ONLY as a labelled negative control).12"""13from __future__ import annotations14import json, re, time, yaml, random15from pathlib import Path16import numpy as np17import mlx.core as mx18from .util import ROOT, read_json, sha256_json, sha256_text, utc_now, read_jsonl19from .records import new_record, write_record, trial_key20from .data.prepare import load_doc, load_key21from .data.synthetic import fmt_amt22from .scoring import score_synthetic23from .prompts import SYSTEM_RECEIVER, receiver_user_text24from .backends.mlx_backend import MLXModel25from .cache.sparse import SparseKVCache, cache_tensor_bytes, dense_kv26from .methods import text as TM2728import platform as _pf29HARDWARE_ID = {"MacBookPro-de-simon-pierre-2.local": "laptop-m5max-48g", "MacStudnpierre2": "m3u96a-m3ultra-96g", "m3u96b": "m3u96b-m3ultra-96g"}.get(_pf.node().split(".")[0] if _pf.node().startswith("MacStud") or _pf.node().startswith("m3u") else _pf.node(), _pf.node())3031def partition_document(doc_text: str, key: dict, k: int, rng: random.Random, other_doc_text: str) -> list[dict]:32 blocks = doc_text.split("\n\n")33 header = blocks[0]; review = blocks[1]34 proj = [b for b in blocks[2:] if b.startswith("Project ")]35 fill = [b for b in blocks[2:] if not b.startswith("Project ")]36 sources = [{"source_id": f"S{i+1}", "blocks": [header]} for i in range(k)]37 for i, b in enumerate(proj):38 sources[i % k]["blocks"].append(b)39 for i, b in enumerate(fill):40 sources[i % k]["blocks"].append(b)41 # overlap with superseded earlier revision: pick a project with >=1 revision42 rev_facts = [f for f in key["facts"] if f["attribute"].startswith("revision_1")]43 conflict = None44 if rev_facts and k >= 2:45 f = rev_facts[0]; code = f["entity"]46 blk = next((b for b in proj if b.startswith(f"Project {code} ")), None)47 if blk:48 # earlier date than the first revision: use the proposed date + 1 day (always earlier than revisions)49 m = re.search(rf"Project {re.escape(code)} was proposed by .*? on (\d+ \w+ \d{{4}})", blk)50 amt = int(re.sub(r"[^\d]", "", f["value"].split()[0])); bogus = fmt_amt(int(amt * 0.93) // 5 * 5)51 cur = f["value"].split()[-1]52 earlier = f"On {m.group(1)} (preliminary) the budget of {code} was revised to {bogus} {cur}; this preliminary entry is superseded by any later-dated entry."53 tgt = (proj.index(blk) + 1) % k54 sources[tgt]["blocks"].append(blk.split(". ")[0] + ". " + earlier)55 conflict = {"code": code, "bogus_amount": bogus, "source": sources[tgt]["source_id"]}56 # irrelevant block from another document57 other_proj = [b for b in other_doc_text.split("\n\n") if b.startswith("Project ")]58 if other_proj:59 sources[-1]["blocks"].append(other_proj[0])60 for s in sources:61 s["text"] = "\n\n".join(s["blocks"])62 return sources, conflict6364def cache_concat_negative_control(receiver: MLXModel, source_texts: list[str]):65 """Prefill each source independently (positions 0..T_i), then concatenate K/V along positions. Overlapping positions66 and no cross-source attention during prefill make this an INVALID composition; measured as a negative control."""67 caches = []68 for t in source_texts:69 ids = receiver.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": "Document:\n" + t + "\n\n"}])70 c, _ = receiver.prefill(ids); caches.append(c)71 out = []72 for layer in range(len(caches[0])):73 ks = [dense_kv(c[layer])[0] for c in caches]; vs = [dense_kv(c[layer])[1] for c in caches]74 k = mx.concatenate(ks, axis=2); v = mx.concatenate(vs, axis=2)75 out.append(SparseKVCache.from_dense(k, v, logical_offset=max(c[layer].offset for c in caches)))76 return out7778def run_fanin(config_path: str, resume: bool = False) -> int:79 cfg = yaml.safe_load(open(config_path))80 mcfg = yaml.safe_load(open(ROOT / "configs/models.yaml"))["models"]81 models = {}82 def get(a):83 if a not in models: models[a] = MLXModel(mcfg[a]["repo"])84 return models[a]85 man = read_json(ROOT / "data/manifests/splits.json")86 items = [it for it in man["items"] if it["family"] == "synthetic" and it["split"] in cfg["splits"] and it.get("target_tokens") in cfg["lengths"]]87 if cfg.get("limit_docs"): items = items[: cfg["limit_docs"]]88 out_path = ROOT / "runs/raw" / f"{cfg['experiment_id']}.jsonl"; log_path = ROOT / "runs/logs" / f"{cfg['experiment_id']}.log"89 done = {r["trial_id"] for r in read_jsonl(out_path) if r["status"] == "success"} if (resume and out_path.exists()) else set()90 protocol_hash = cfg.get("protocol_hash", "unfrozen"); config_hash = sha256_json(cfg)91 def log(m):92 line = f"{utc_now()} {m}"; print(line, flush=True)93 with open(log_path, "a") as f: f.write(line + "\n")94 QTYPES = set(cfg.get("question_types", ["multi_hop", "aggregation", "revision", "temporal", "unanswerable"]))95 agg = get(cfg["aggregator"]); sender = get(cfg["sender"])96 for idx, it in enumerate(items):97 doc = load_doc(it["document_id"], "synthetic"); key = load_key(it["document_id"], "synthetic")98 other = load_doc(items[(idx + 1) % len(items)]["document_id"], "synthetic")99 qkeys = {q["question_id"]: q for q in key["questions"]}100 qs = [q for q in doc["questions"] if q["qtype"] in QTYPES]101 for k in cfg["fan_in"]:102 rng = random.Random(it["seed"] * 7 + k)103 sources, conflict = partition_document(doc["text"], key, k, rng, other["text"])104 for method in cfg["methods"]:105 cid = f"{method}|k={k}"106 tks = {q["question_id"]: trial_key(cfg["experiment_id"], it["document_id"], q["question_id"], method, {"k": k, "budget": cfg["budget_tokens"]}, cfg["sender"], cfg["aggregator"], "blind", protocol_hash, 0, "fanin") for q in qs}107 if all(t in done for t in tks.values()): continue108 t0 = time.perf_counter(); info = {"k": k, "conflict": conflict}109 try:110 cache = None; material = None; per_q = False111 if method == "full_context":112 material = "\n\n".join(f"[Source {s['source_id']}]\n{s['text']}" for s in sources); label = "Documents from all sources"113 elif method == "concat_A4":114 msgs = [TM.a4_summary(s["text"], cfg["budget_tokens"], sender, None) for s in sources]115 material = "\n\n".join(f"[Source {s['source_id']} summary]\n{m.text}" for s, m in zip(sources, msgs)); label = "Summaries from several sources"116 info["sender_tokens"] = sum(m.sender_usage.get("sender_output_tokens", 0) for m in msgs)117 elif method == "merge_A5":118 msgs = [TM.a5_structured(s["text"], cfg["budget_tokens"], sender, None) for s in sources]119 facts = []120 for s, m in zip(sources, msgs):121 for f in json.loads(m.text).get("facts", []):122 facts.append({"source": s["source_id"], **{kk: vv for kk, vv in f.items() if kk != "span"}})123 material = json.dumps({"facts": facts, "rule": "later-dated entries supersede earlier ones"}, ensure_ascii=False); label = "Merged structured evidence from several sources (JSON)"124 info["n_facts"] = len(facts)125 elif method == "shared_memory_A6":126 msgs = [TM.a6_references(s["text"], cfg["budget_tokens"], sender, None) for s in sources]127 per_q = True; label = "Reference index from several sources plus retrieved passages"128 elif method == "cache_concat_negative_control":129 cache = cache_concat_negative_control(agg, [s["text"] for s in sources]); label = None130 info["cache_bytes"] = cache_tensor_bytes(cache)131 else:132 raise NotImplementedError(method)133 build_ms = (time.perf_counter() - t0) * 1000134 except Exception as e:135 log(f"BUILD FAIL {cid} {it['document_id']}: {type(e).__name__}: {str(e)[:200]}"); continue136 for q in qs:137 tk = tks[q["question_id"]]138 if tk in done: continue139 rec = new_record(experiment_id=cfg["experiment_id"], trial_id=tk, run_group_id=f"{cfg['experiment_id']}:{cid}", protocol_hash=protocol_hash, config_hash=config_hash,140 task_family="synthetic", document_id=it["document_id"], document_hash=it["document_hash"], question_id=q["question_id"], split=it["split"], cluster_id=it["cluster_id"],141 sender_model_revision=f"{sender.info.repo_id}@{sender.info.revision}", receiver_model_revision=f"{agg.info.repo_id}@{agg.info.revision}", tokenizer_hashes={"receiver": agg.info.tokenizer_hash},142 backend_versions={"mlx": mx.__version__, "mlx_lm": agg.info.mlx_lm_version}, hardware_id=HARDWARE_ID, device_state={}, weight_precision="bf16", method_id=method,143 implementation_fidelity="custom" if method != "cache_concat_negative_control" else "n/a", compression_config={"k": k, "budget_tokens": cfg["budget_tokens"], **{kk: vv for kk, vv in info.items() if kk != "conflict"}},144 question_access="blind", topology=f"fanin_k{k}", handoff_index=1, reuse_count=len(qs), cache_hit=False, cache_compatibility_hash=None, random_seed=0,145 source_tokens_sender=None, source_tokens_receiver=len(agg.encode(doc["text"])), fetched_bytes=0, side_information_bytes=0, timing_spans={"build_ms": build_ms},146 input_usage={}, output_usage={}, provider_cached_usage=None, process_memory=None, backend_memory=None, peak_prefill_memory=None, peak_decode_memory=None, mapper_memory=None,147 cost_with_source=None, evidence_artifact={"conflict": conflict, "question_type": q["qtype"]}, reference_key_id=f"keys/synthetic/{it['document_id']}.json", status="implementation_error",148 error_code=None, log_artifact=str(log_path.relative_to(ROOT)), payload_hash=None, output_hash=None, message_utf8_bytes=None, serialized_payload_bytes=None, cache_tensor_bytes=info.get("cache_bytes"), query_tokens=None, message_tokens=None)149 try:150 t1 = time.perf_counter()151 if cache is not None:152 from .cache.sparse import copy_cache153 # query tokens = the fresh tail of a full prompt built on the last source (identical tail for every source)154 full_ids = agg.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": receiver_user_text("Document", sources[-1]["text"], q["question"], "synthetic")}])155 pre_ids = agg.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": "Document:\n" + sources[-1]["text"] + "\n\n"}])156 n = 0157 while n < min(len(pre_ids), len(full_ids)) and pre_ids[n] == full_ids[n]: n += 1158 q_ids = full_ids[n:]159 g = agg.generate_from_cache(copy_cache(cache), q_ids, max_tokens=48)160 rec.update({"serialized_payload_bytes": info["cache_bytes"], "message_utf8_bytes": sum(len(s["text"].encode()) for s in sources), "query_tokens": len(q_ids)})161 else:162 mat = material163 if per_q:164 pieces = []; fb = 0; lk = 0165 for m in msgs:166 mt, ri = TM.a6_retrieve(m, q["question"], cfg.get("max_fetch_bytes", 1500), cfg.get("max_lookups", 3)); pieces.append(mt); fb += ri["fetched_bytes"]; lk += ri["lookups"]167 mat = "\n\n".join(pieces); rec["fetched_bytes"] = fb; rec["input_usage"] = {"lookups": lk}168 ids = agg.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": receiver_user_text(label, mat, q["question"], "synthetic")}])169 g = agg.generate_text(ids, max_tokens=48)170 rec.update({"serialized_payload_bytes": len(mat.encode("utf-8")), "message_utf8_bytes": len(mat.encode("utf-8")), "query_tokens": len(ids), "message_tokens": len(agg.encode(mat)), "payload_hash": sha256_text(mat)})171 rec.update({"prediction_artifact": g["text"], "output_hash": sha256_text(g["text"]), "ttft_ms": g["ttft_ms"], "total_wall_ms": (time.perf_counter() - t1) * 1000 + build_ms,172 "output_usage": {"output_tokens": g["output_tokens"]}, "task_scores": score_synthetic(g["text"], qkeys[q["question_id"]]), "status": "success" if g["text"].strip() else "invalid_output"})173 rec["timing_spans"].update({"ttft_ms": g["ttft_ms"], "decode_ms": g["decode_ms"], "receiver_prefill_ms": g.get("prefill_ms", 0.0)})174 except Exception as e:175 rec["error_code"] = f"{type(e).__name__}: {str(e)[:200]}"176 write_record(rec, out_path)177 sc = rec["task_scores"] or {}178 log(f"{cid:>36s} {it['document_id']} {q['question_id']:>18s} {rec['status']:>8s} corr={sc.get('correct')} bytes={rec['serialized_payload_bytes']} ans={json.dumps((rec['prediction_artifact'] or '')[:50])}")179 mx.clear_cache()180 log("fanin end")181 return 0182