muster-memory 0.1.0__py3-none-any.whl

muster/compat/ce_cache.py

@@ -0,0 +1,47 @@
+"""The harness's reranker score cache, read-only (harness/lme_pf29.py:1774-1848).
+
+A file is named `<corpus signature>__<knob digest>.json`. The corpus signature hashes the exact unit list, so a
+cut that names an existing file has produced, byte for byte, the units the published run scored.
+"""
+import hashlib
+import json
+import os
+
+RERANKER_NAME = "BAAI/bge-reranker-v2-m3"
+# The knob string of every final leg except FactConsolidation (len_legs.json env: CEML=2048 CEBS=16 CE_SUBMAX=0).
+KNOBS_FINAL = "CEML=2048|CEBS=16|CE_SUBMAX=0|CE_SUB_TOK=|CE_SLIDE=|name=" + RERANKER_NAME
+
+
+def ce_sig(units, name=RERANKER_NAME):
+    h = hashlib.sha1()
+    h.update(("ce1|%s|%d\n" % (name, len(units))).encode())
+    for t in units:
+        h.update(hashlib.sha1(t.encode("utf-8", "replace")).digest())
+    return h.hexdigest()[:16]
+
+
+def knob_digest(knobs=KNOBS_FINAL):
+    return hashlib.sha1(knobs.encode()).hexdigest()[:8]
+
+
+def qkey(query):
+    return hashlib.sha1(query.encode("utf-8", "replace")).hexdigest()[:16]
+
+
+class LegacyCECache:
+    """{question key: {unit index: score}} for one corpus, from a directory of harness cache files."""
+
+    def __init__(self, directory, *, knobs=KNOBS_FINAL):
+        self.directory, self.knobs, self.suffix = directory, knobs, knob_digest(knobs)
+
+    def path_for(self, units):
+        return os.path.join(self.directory, "%s__%s.json" % (ce_sig(units), self.suffix))
+
+    def has(self, units):
+        return os.path.exists(self.path_for(units))
+
+    def load(self, units):
+        p = self.path_for(units)
+        if not os.path.exists(p):
+            return {}
+        return json.load(open(p))

muster/compat/env.py

@@ -0,0 +1,39 @@
+"""The environment exact mode needs, and a record of the one a run used.
+
+Exact mode iterates sets the way the paper's harness did, which is reproducible only with string hashing fixed
+(PYTHONHASHSEED=0 at interpreter start; `python -I` and `-E` ignore the variable). The published equality was
+measured with the versions in EXACT_VERSIONS on CPython 3.12.13.
+"""
+import importlib.metadata
+import platform
+import sys
+
+from ..errors import EnvironmentRefused
+
+EXACT_PYTHON = "3.12.13"
+EXACT_VERSIONS = {"tiktoken": "0.13.0", "numpy": "2.4.4", "torch": "2.11.0", "sentence-transformers": "5.3.0",
+                  "transformers": "5.5.0", "tokenizers": "0.22.2", "nltk": "3.10.0", "rank-bm25": "0.2.2"}
+
+
+def versions():
+    out = {"python": platform.python_version(), "hash_randomization": sys.flags.hash_randomization}
+    for k in EXACT_VERSIONS:
+        try:
+            out[k] = importlib.metadata.version(k)
+        except importlib.metadata.PackageNotFoundError:
+            out[k] = None
+    return out
+
+
+def check_environment(*, strict_versions=False):
+    """Raise EnvironmentRefused if exact mode cannot reproduce the paper here; return the mismatches otherwise."""
+    if sys.flags.hash_randomization:
+        raise EnvironmentRefused("exact mode needs string hashing fixed: start Python as "
+                                 "`PYTHONHASHSEED=0 python ...` (not with -I or -E)")
+    v = versions()
+    bad = {k: (v.get(k), want) for k, want in EXACT_VERSIONS.items() if v.get(k) != want}
+    if v["python"] != EXACT_PYTHON:
+        bad["python"] = (v["python"], EXACT_PYTHON)
+    if bad and strict_versions:
+        raise EnvironmentRefused("exact mode was verified with other versions: %s" % bad)
+    return bad

muster/compat/intermediates.py

@@ -0,0 +1,140 @@
+"""The harness's intermediate renders, from which the published R4 contexts were made (exact mode only).
+
+The paper's R4 contexts were produced in two steps: the harness rendered the selected units with an intermediate
+grouping (bgroup for prose stores, seg3 for BEAM and LoCoMo), and render_rule.py then recovered the unit ids from that
+text and rendered one block per parent. Where a store holds several units with identical text, the recovery can name
+another copy than the one selected, and the published context follows the copy. Reproducing a published context byte
+for byte therefore needs both steps. Lifted from harness/lme_pf29.py:271-333 (bgroup) and 525-715 (seg3, HMARK off).
+"""
+import math
+
+
+def rank_pos(order):
+    """First position of each unit in a ranking."""
+    pos = {}
+    for n, i in enumerate(order):
+        pos.setdefault(i, n)
+    return pos
+
+
+def render_bgroup(fed, txt, pos, far, parent, ents, hubs, cap):
+    fed = sorted(fed)
+    lim = max(1, cap if cap > 0 else 10)
+    if len(fed) <= lim:
+        return "\n".join("Memory %d:\n%s" % (n + 1, txt(i))
+                         for n, i in enumerate(sorted(fed, key=lambda j: (pos.get(j, far), j))))
+    E = {i: (ents[i] - hubs) if i < len(ents) else set() for i in fed}
+    _ov = [len(E[fed[a]] & E[fed[b]]) for a in range(len(fed)) for b in range(a + 1, len(fed))]
+    PB = (max(_ov) if _ov else 0) + 1
+    groups = {i: [i] for i in fed}
+
+    def sim(a, b):
+        s = 0
+        for x in groups[a]:
+            for y in groups[b]:
+                s += len(E[x] & E[y]) + (PB if parent[x] == parent[y] else 0)
+        return s / float(len(groups[a]) * len(groups[b]))
+
+    size_cap = int(math.ceil(len(fed) / float(lim)))
+    while len(groups) > lim:
+        ks = sorted(groups)
+        pick, best = None, -1.0
+        for ia in range(len(ks)):
+            for ib in range(ia + 1, len(ks)):
+                a, b = ks[ia], ks[ib]
+                if len(groups[a]) + len(groups[b]) > size_cap:
+                    continue
+                s = sim(a, b)
+                if s > best:
+                    pick, best = (a, b), s
+        if pick is None:
+            size_cap += 1
+            continue
+        a, b = pick
+        groups[a] = groups[a] + groups[b]
+        del groups[b]
+    out = sorted(groups.values(), key=lambda m: (min(pos.get(i, far) for i in m), m[0]))
+    return "\n".join("Memory %d:\n%s" % (n + 1, "\n".join(txt(i) for i in sorted(m))) for n, m in enumerate(out))
+
+
+def render_seg3(fed, txt, parent, ents, hubs, cap):
+    """seg3 (doc_first) with HMARK off: an ordered segmentation of the sorted fed units, merged by priority levels."""
+    fed = sorted(fed)
+    P = len({parent[i] for i in fed})
+    lim = min(len(fed), max(-cap, P)) if cap < 0 else max(1, cap if cap > 0 else 10)
+    if len(fed) <= lim:
+        return "\n".join("Memory %d:\n%s" % (n + 1, txt(i)) for n, i in enumerate(fed))
+    E = {i: (ents[i] - hubs) if i < len(ents) else set() for i in fed}
+    D = {}
+    for i, p in enumerate(parent):
+        if i < len(ents):
+            D.setdefault(p, set()).update(ents[i])
+    D = {p: (s - hubs) for p, s in D.items()}
+    blocks = [[i] for i in fed]
+
+    def pars(b):
+        return {parent[i] for i in b}
+
+    def jac(sa, sb):
+        u = len(sa | sb)
+        return (len(sa & sb) / float(u)) if u else 0.0
+
+    def b1(a, b):
+        s = 0
+        for x in a:
+            for y in b:
+                s += len(E[x] & E[y])
+        return s / float(len(a) * len(b))
+
+    def b2(a, b):
+        return jac(set().union(*[E[i] for i in a]), set().union(*[E[i] for i in b]))
+
+    def b3(a, b):
+        sa = set().union(*[D.get(p, set()) for p in pars(a)]) if a else set()
+        sb = set().union(*[D.get(p, set()) for p in pars(b)]) if b else set()
+        return jac(sa, sb)
+
+    _sing = [(x, y) for ix, x in enumerate(fed) for y in fed[ix + 1:] if parent[x] != parent[y]]
+    tau_i = (sum(b1([x], [y]) for x, y in _sing) / len(_sing)) if _sing else 0.0
+    _pl = sorted({parent[i] for i in fed}, key=lambda p: min(i for i in fed if parent[i] == p))
+    _grp = {p: [i for i in fed if parent[i] == p] for p in _pl}
+    _pp = [(a, b) for ia, a in enumerate(_pl) for b in _pl[ia + 1:]]
+    tau_b = (sum(b2(_grp[a], _grp[b]) for a, b in _pp) / len(_pp)) if _pp else 0.0
+    tau_c = (sum(b3(_grp[a], _grp[b]) for a, b in _pp) / len(_pp)) if _pp else 0.0
+    LEVELS = [(1, b1, tau_i, True), (1, b1, 0.0, True), (2, b2, tau_b, False), (2, b2, 0.0, False),
+              (3, b3, tau_c, False), (3, b3, 0.0, False), (4, None, None, False)]
+    size_cap = int(math.ceil(len(fed) / float(lim)))
+    while len(blocks) > lim:
+        got = False
+        for lv, bond, thr, same_only in LEVELS:
+            if lv == 4:
+                pool = [j for j in range(len(blocks) - 1) if len(blocks[j]) + len(blocks[j + 1]) <= size_cap]
+                if not pool:
+                    continue
+                k = min(pool, key=lambda j: len(blocks[j]) + len(blocks[j + 1]))
+                blocks[k] = blocks[k] + blocks[k + 1]
+                del blocks[k + 1]
+                got = True
+                break
+            cand = []
+            for j in range(len(blocks) - 1):
+                A, B = blocks[j], blocks[j + 1]
+                one = (len(pars(A)) == 1 and pars(A) == pars(B))
+                if same_only != one:
+                    continue
+                if not one and len(A) + len(B) > size_cap:
+                    continue
+                s = bond(A, B)
+                if s < thr or (s <= 0 and not (lv == 1 and thr == 0.0)):
+                    continue
+                cand.append((s, j))
+            if not cand:
+                continue
+            s, j = max(cand, key=lambda t: (t[0], -t[1]))
+            blocks[j] = blocks[j] + blocks[j + 1]
+            del blocks[j + 1]
+            got = True
+            break
+        if not got:
+            size_cap += 1
+    return "\n".join("Memory %d:\n%s" % (n + 1, "\n".join(txt(i) for i in m)) for n, m in enumerate(blocks))

muster/compat/recover.py

@@ -0,0 +1,75 @@
+"""Recover the selected unit ids from a rendered harness context (harness/render_parent.py:13-43,
+harness/rerender_nomerge.py:53-55, harness/render_rule.py:28-29), so R4 can be rendered from what a published run fed.
+"""
+import collections
+import re
+
+HDR = re.compile(r"(?m)^Memory \d+:\n")
+
+
+def index80(units):
+    idx = {}
+    for i, s in enumerate(units):
+        idx.setdefault(s[:80], []).append(i)
+    return idx
+
+
+def blocks_of(context):
+    """The block bodies of a "Memory n:" context, joiners stripped (render_rule.py:28-29)."""
+    parts = HDR.split(context)[1:]
+    return [b[:-1] if (n < len(parts) - 1 and b.endswith("\n")) else b for n, b in enumerate(parts)]
+
+
+def decompose(block, units, idx):
+    """Unit ids of one block; among same-text units, the first one after the block's previous unit."""
+    ids, p = [], 0
+    while p < len(block):
+        cand = [i for i in idx.get(block[p:p + 80], []) if block.startswith(units[i], p)
+                and (p + len(units[i]) == len(block) or block[p + len(units[i])] == "\n")]
+        if not cand:
+            return None
+        L = max(len(units[j]) for j in cand)
+        cand = [j for j in cand if len(units[j]) == L]
+        after = [j for j in cand if not ids or j > ids[-1]]
+        i = min(after) if after else min(cand)
+        ids.append(i)
+        p += len(units[i])
+        if p < len(block):
+            p += 1
+    return ids
+
+
+def resolve(dec, units, parent):
+    """Where one id was given to two places (a repeated passage), move the extra one to an unused copy of the same
+    text, preferably in the parent of its block-mates."""
+    used = collections.Counter(i for d in dec for i in d)
+    copies = collections.defaultdict(list)
+    for j, s in enumerate(units):
+        copies[s].append(j)
+    for d in dec:
+        for p, i in enumerate(d):
+            if used[i] < 2:
+                continue
+            mates = {parent[x] for x in d if x != i}
+            free = [j for j in copies[units[i]] if used[j] == 0]
+            alt = [j for j in free if mates and parent[j] in mates] or free
+            if alt:
+                used[i] -= 1
+                d[p] = alt[0]
+                used[alt[0]] += 1
+    return dec
+
+
+def recover_ids(context, units, parent, idx=None):
+    """The ids a harness context holds, or raise ValueError (a header-less context, a failed decomposition)."""
+    if not HDR.search(context):
+        raise ValueError("context has no 'Memory n:' header; nothing to recover")
+    idx = idx if idx is not None else index80(units)
+    dec = [decompose(b, units, idx) for b in blocks_of(context)]
+    if any(d is None for d in dec):
+        raise ValueError("unit decomposition failed")
+    dec = resolve(dec, units, parent)
+    ids = [i for d in dec for i in d]
+    if len(set(ids)) != len(ids):
+        raise ValueError("a unit id is used twice")
+    return ids

muster/compat/verify.py

@@ -0,0 +1,142 @@
+"""Check the library against the paper without the paper's machine: tables of hashes, no benchmark text.
+
+  muster verify cuts --data DIR      every row of the 24 final legs: unit count, unit-list hash, parent-label hash
+
+DIR holds the benchmark files at the paths the legs name (data/..., data_all/data/..., beam/...), each checked against
+gate/data.sha256 first. The table (gate/cut_table.json) was written by tools/make_cut_table.py on the experiment machine
+from the same library code, after the cut had been checked row by row against the harness (tests/studio/).
+"""
+import hashlib
+import json
+import os
+
+from .. import presets
+from ..cut import cut
+from .ce_cache import ce_sig
+
+
+def parent_hash(parent):
+    h = hashlib.sha1()
+    for p in parent:
+        h.update(("%s:%s\n" % (type(p).__name__, p)).encode())
+    return h.hexdigest()[:16]
+
+
+def text_hash(text):
+    return hashlib.sha256(text.encode("utf-8")).hexdigest()
+
+
+def row_signature(text, plan):
+    c = cut(text, plan)
+    return {"units": len(c.units), "unit_sig": ce_sig(list(c.units)), "parent_sig": parent_hash(c.parent),
+            "parents": len(set(c.parent)), "text_sha256": text_hash(text)}
+
+
+def _sha256(path):
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for b in iter(lambda: f.read(1 << 20), b""):
+            h.update(b)
+    return h.hexdigest()
+
+
+def load_data_hashes(path):
+    out = {}
+    for line in open(path):
+        if line.strip() and not line.startswith("#"):
+            sha, rel = line.split()
+            out[rel] = sha
+    return out
+
+
+def verify_cuts(data_dir, table_path, hashes_path=None, *, legs=None, out=print):
+    """Each row's store text is checked against the table's text hash before it is cut, so a file rebuilt with
+    another parquet writer still verifies, and a wrong file is reported, never compared."""
+    import pandas as pd
+    table = json.load(open(table_path))
+    ok = bad = skipped = 0
+    for tag, rows in sorted(table["legs"].items()):
+        if legs and tag not in legs:
+            continue
+        leg = presets.legs()[tag]
+        path = os.path.join(data_dir, leg["parq"])
+        if not os.path.exists(path):
+            out("SKIP %-16s %s not found" % (tag, leg["parq"]))
+            skipped += len(rows)
+            continue
+        df = pd.read_parquet(path)
+        plan = presets.get(tag)
+        n_wrong = 0
+        for r, want in rows.items():
+            text = str(df["context"].iloc[int(r)]) if int(r) < len(df) else None
+            if text is None or ("text_sha256" in want and text_hash(text) != want["text_sha256"]):
+                n_wrong += 1
+                skipped += 1
+                continue
+            got = row_signature(text, plan)
+            if "text_sha256" not in want:
+                got.pop("text_sha256")
+            same = got == want
+            ok += same
+            bad += not same
+            if not same:
+                out("DIFF %-16s row %s: %s != %s" % (tag, r, got, want))
+        out("%-16s %d rows checked%s" % (tag, len(rows) - n_wrong,
+                                          (", %d skipped: store text differs from the published one" % n_wrong) if n_wrong else ""))
+    out("cut verification: %d equal, %d different, %d skipped (missing or different data)" % (ok, bad, skipped))
+    return bad == 0 and ok > 0
+
+
+def verify_contexts(data_dir, ctx_table_path, cut_table_path, *, legs=None, cache_dirs=(), models=None, nq=None,
+                    out=print):
+    """Rebuild each leg's stores under its preset, search every question in exact mode and compare the context hash
+    with gate/ctx_table.jsonl. With the published reranker caches (cache_dirs) the comparison is exact on CPU;
+    without them, scores are computed live and near-equal scores may come out in another order."""
+    import pandas as pd
+    from ..memory import build
+    from ..models import Models
+    from .ce_cache import LegacyCECache
+    want = {}
+    for line in open(ctx_table_path):
+        r = json.loads(line)
+        want.setdefault(r["leg"], {})[(r["row"], r["qi"])] = r["ctxhash"]
+    texts = {tag: {int(r): v.get("text_sha256") for r, v in rows.items()}
+             for tag, rows in json.load(open(cut_table_path))["legs"].items()}
+    total = equal = 0
+    for tag in sorted(want):
+        if legs and tag not in legs:
+            continue
+        leg = presets.legs()[tag]
+        path = os.path.join(data_dir, leg["parq"])
+        if not os.path.exists(path):
+            out("SKIP %-16s %s not found" % (tag, leg["parq"]))
+            continue
+        plan = presets.get(tag)
+        m = models or Models(device="cpu", encoder_spec=plan.encoder, reranker_spec=plan.reranker)
+        df = pd.read_parquet(path)
+        t = e = 0
+        for row in leg["rows"]:
+            keys = sorted(k for k in want[tag] if k[0] == row)[:nq] if nq else sorted(k for k in want[tag] if k[0] == row)
+            if not keys:
+                continue
+            text = str(df["context"].iloc[row])
+            if texts.get(tag, {}).get(row) and text_hash(text) != texts[tag][row]:
+                out("SKIP %-16s row %d: store text differs from the published one" % (tag, row))
+                continue
+            mem = build(text, plan, models=m)
+            cache = {}
+            for d in cache_dirs:
+                c = LegacyCECache(d, knobs=plan.compat["knobs"])
+                if c.has(list(mem.units)):
+                    cache = c.load(list(mem.units))
+                    break
+            qs = list(df["questions"].iloc[row])
+            for k in keys:
+                ctx = mem.context(str(qs[k[1]]), cache=cache, models=m)
+                t += 1
+                e += ctx.ctxhash == want[tag][k]
+        total += t
+        equal += e
+        out("%-16s %5d questions, %5d contexts equal" % (tag, t, e))
+    out("context verification: %d of %d equal" % (equal, total))
+    return total > 0 and equal == total