"""E5: fan-in / composition with conflicting evidence. The synthetic document is partitioned into K source agents (project blocks round-robin). Overlap: one project block is duplicated into a second source with an *earlier* (superseded) revision sentence, so the ledger's precedence rule still yields the original answers but a naive merge sees a conflict. One source additionally receives an irrelevant block from another document. The aggregator answers multi-hop/aggregation/revision questions that need several sources. Aggregation methods: concat_A4 (concatenated prose summaries), merge_A5 (merged structured records), shared_memory_A6 (reference index per source + metered retrieval by the aggregator), full_context (reference: all source texts concatenated), and cache_concat_negative_control (KV caches of independently prefilled sources concatenated along the position axis — positions overlap and no cross-source context exists; reported ONLY as a labelled negative control). """ from __future__ import annotations import json, re, time, yaml, random from pathlib import Path import numpy as np import mlx.core as mx from .util import ROOT, read_json, sha256_json, sha256_text, utc_now, read_jsonl from .records import new_record, write_record, trial_key from .data.prepare import load_doc, load_key from .data.synthetic import fmt_amt from .scoring import score_synthetic from .prompts import SYSTEM_RECEIVER, receiver_user_text from .backends.mlx_backend import MLXModel from .cache.sparse import SparseKVCache, cache_tensor_bytes, dense_kv from .methods import text as TM import platform as _pf HARDWARE_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()) def partition_document(doc_text: str, key: dict, k: int, rng: random.Random, other_doc_text: str) -> list[dict]: blocks = doc_text.split("\n\n") header = blocks[0]; review = blocks[1] proj = [b for b in blocks[2:] if b.startswith("Project ")] fill = [b for b in blocks[2:] if not b.startswith("Project ")] sources = [{"source_id": f"S{i+1}", "blocks": [header]} for i in range(k)] for i, b in enumerate(proj): sources[i % k]["blocks"].append(b) for i, b in enumerate(fill): sources[i % k]["blocks"].append(b) # overlap with superseded earlier revision: pick a project with >=1 revision rev_facts = [f for f in key["facts"] if f["attribute"].startswith("revision_1")] conflict = None if rev_facts and k >= 2: f = rev_facts[0]; code = f["entity"] blk = next((b for b in proj if b.startswith(f"Project {code} ")), None) if blk: # earlier date than the first revision: use the proposed date + 1 day (always earlier than revisions) m = re.search(rf"Project {re.escape(code)} was proposed by .*? on (\d+ \w+ \d{{4}})", blk) amt = int(re.sub(r"[^\d]", "", f["value"].split()[0])); bogus = fmt_amt(int(amt * 0.93) // 5 * 5) cur = f["value"].split()[-1] 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." tgt = (proj.index(blk) + 1) % k sources[tgt]["blocks"].append(blk.split(". ")[0] + ". " + earlier) conflict = {"code": code, "bogus_amount": bogus, "source": sources[tgt]["source_id"]} # irrelevant block from another document other_proj = [b for b in other_doc_text.split("\n\n") if b.startswith("Project ")] if other_proj: sources[-1]["blocks"].append(other_proj[0]) for s in sources: s["text"] = "\n\n".join(s["blocks"]) return sources, conflict def cache_concat_negative_control(receiver: MLXModel, source_texts: list[str]): """Prefill each source independently (positions 0..T_i), then concatenate K/V along positions. Overlapping positions and no cross-source attention during prefill make this an INVALID composition; measured as a negative control.""" caches = [] for t in source_texts: ids = receiver.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": "Document:\n" + t + "\n\n"}]) c, _ = receiver.prefill(ids); caches.append(c) out = [] for layer in range(len(caches[0])): ks = [dense_kv(c[layer])[0] for c in caches]; vs = [dense_kv(c[layer])[1] for c in caches] k = mx.concatenate(ks, axis=2); v = mx.concatenate(vs, axis=2) out.append(SparseKVCache.from_dense(k, v, logical_offset=max(c[layer].offset for c in caches))) return out def run_fanin(config_path: str, resume: bool = False) -> int: cfg = yaml.safe_load(open(config_path)) mcfg = yaml.safe_load(open(ROOT / "configs/models.yaml"))["models"] models = {} def get(a): if a not in models: models[a] = MLXModel(mcfg[a]["repo"]) return models[a] man = read_json(ROOT / "data/manifests/splits.json") 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"]] if cfg.get("limit_docs"): items = items[: cfg["limit_docs"]] out_path = ROOT / "runs/raw" / f"{cfg['experiment_id']}.jsonl"; log_path = ROOT / "runs/logs" / f"{cfg['experiment_id']}.log" done = {r["trial_id"] for r in read_jsonl(out_path) if r["status"] == "success"} if (resume and out_path.exists()) else set() protocol_hash = cfg.get("protocol_hash", "unfrozen"); config_hash = sha256_json(cfg) def log(m): line = f"{utc_now()} {m}"; print(line, flush=True) with open(log_path, "a") as f: f.write(line + "\n") QTYPES = set(cfg.get("question_types", ["multi_hop", "aggregation", "revision", "temporal", "unanswerable"])) agg = get(cfg["aggregator"]); sender = get(cfg["sender"]) for idx, it in enumerate(items): doc = load_doc(it["document_id"], "synthetic"); key = load_key(it["document_id"], "synthetic") other = load_doc(items[(idx + 1) % len(items)]["document_id"], "synthetic") qkeys = {q["question_id"]: q for q in key["questions"]} qs = [q for q in doc["questions"] if q["qtype"] in QTYPES] for k in cfg["fan_in"]: rng = random.Random(it["seed"] * 7 + k) sources, conflict = partition_document(doc["text"], key, k, rng, other["text"]) for method in cfg["methods"]: cid = f"{method}|k={k}" 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} if all(t in done for t in tks.values()): continue t0 = time.perf_counter(); info = {"k": k, "conflict": conflict} try: cache = None; material = None; per_q = False if method == "full_context": material = "\n\n".join(f"[Source {s['source_id']}]\n{s['text']}" for s in sources); label = "Documents from all sources" elif method == "concat_A4": msgs = [TM.a4_summary(s["text"], cfg["budget_tokens"], sender, None) for s in sources] 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" info["sender_tokens"] = sum(m.sender_usage.get("sender_output_tokens", 0) for m in msgs) elif method == "merge_A5": msgs = [TM.a5_structured(s["text"], cfg["budget_tokens"], sender, None) for s in sources] facts = [] for s, m in zip(sources, msgs): for f in json.loads(m.text).get("facts", []): facts.append({"source": s["source_id"], **{kk: vv for kk, vv in f.items() if kk != "span"}}) material = json.dumps({"facts": facts, "rule": "later-dated entries supersede earlier ones"}, ensure_ascii=False); label = "Merged structured evidence from several sources (JSON)" info["n_facts"] = len(facts) elif method == "shared_memory_A6": msgs = [TM.a6_references(s["text"], cfg["budget_tokens"], sender, None) for s in sources] per_q = True; label = "Reference index from several sources plus retrieved passages" elif method == "cache_concat_negative_control": cache = cache_concat_negative_control(agg, [s["text"] for s in sources]); label = None info["cache_bytes"] = cache_tensor_bytes(cache) else: raise NotImplementedError(method) build_ms = (time.perf_counter() - t0) * 1000 except Exception as e: log(f"BUILD FAIL {cid} {it['document_id']}: {type(e).__name__}: {str(e)[:200]}"); continue for q in qs: tk = tks[q["question_id"]] if tk in done: continue 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, 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"], 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}, backend_versions={"mlx": mx.__version__, "mlx_lm": agg.info.mlx_lm_version}, hardware_id=HARDWARE_ID, device_state={}, weight_precision="bf16", method_id=method, 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"}}, 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, 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}, 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, 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", 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) try: t1 = time.perf_counter() if cache is not None: from .cache.sparse import copy_cache # query tokens = the fresh tail of a full prompt built on the last source (identical tail for every source) full_ids = agg.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": receiver_user_text("Document", sources[-1]["text"], q["question"], "synthetic")}]) pre_ids = agg.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": "Document:\n" + sources[-1]["text"] + "\n\n"}]) n = 0 while n < min(len(pre_ids), len(full_ids)) and pre_ids[n] == full_ids[n]: n += 1 q_ids = full_ids[n:] g = agg.generate_from_cache(copy_cache(cache), q_ids, max_tokens=48) 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)}) else: mat = material if per_q: pieces = []; fb = 0; lk = 0 for m in msgs: 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"] mat = "\n\n".join(pieces); rec["fetched_bytes"] = fb; rec["input_usage"] = {"lookups": lk} ids = agg.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": receiver_user_text(label, mat, q["question"], "synthetic")}]) g = agg.generate_text(ids, max_tokens=48) 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)}) 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, "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"}) rec["timing_spans"].update({"ttft_ms": g["ttft_ms"], "decode_ms": g["decode_ms"], "receiver_prefill_ms": g.get("prefill_ms", 0.0)}) except Exception as e: rec["error_code"] = f"{type(e).__name__}: {str(e)[:200]}" write_record(rec, out_path) sc = rec["task_scores"] or {} 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])}") mx.clear_cache() log("fanin end") return 0