Upload scripts/longeval_gen.py with huggingface_hub

scripts/longeval_gen.py

@@ -184,6 +184,219 @@ def t_val_backward(rng, depth):
     return [(s, depth)], q, a, [k]
 
 
+
+# ---------------------------------------------------------------------------
+# RULER PARITY ADDITIONS (2026-09-17)
+#
+# RULER is 13 tasks in 4 categories: 8 NIAH retrieval variants, 1 multi-hop
+# tracing, 2 aggregation, 2 QA. The six tasks above cover 4 of the 8 NIAH slots
+# plus VT. Everything below fills the rest.
+#
+# The NIAH variants differ only in NEEDLE TYPE (7-digit number vs UUID) and
+# HAYSTACK TYPE (real text vs a repeated noise sentence), so they are
+# parameterisations rather than new mechanisms -- which is exactly how RULER
+# defines them.
+#
+# CWE and FWE are structurally different: the body IS the data, not a haystack
+# with needles planted in it. They are listed in BODY_TASKS and build_one takes
+# a separate path for them.
+# ---------------------------------------------------------------------------
+
+def rnd_uuid(rng):
+    h = "0123456789abcdef"
+    f = lambda n: "".join(rng.choice(h) for _ in range(n))
+    return f"{f(8)}-{f(4)}-{f(4)}-{f(4)}-{f(12)}"
+
+
+# Tasks whose haystack must be the repeated-noise filler rather than real text.
+NOISE_TASKS = {"niah_single_1"}
+# Tasks that generate their own body instead of planting needles in a haystack.
+BODY_TASKS = {"cwe", "fwe", "qa_1", "qa_2"}
+
+
+def t_niah_single_1(rng, depth):
+    """S-NIAH-1: word-number needle in REPEATED NOISE (see NOISE_TASKS)."""
+    return t_niah_single(rng, depth)
+
+
+def t_niah_single_3(rng, depth):
+    """S-NIAH-3: word-UUID needle. UUIDs are far longer and less token-friendly
+    than 7-digit numbers, so this is the hardest single-needle variant."""
+    k, v = rng.choice(WORDS), rnd_uuid(rng)
+    s = f" The special magic uuid for {k} is {v}. "
+    q = f"\nQuestion: What is the special magic uuid for {k}?\n"
+    a = f"Answer: The special magic uuid for {k} is"
+    return [(s, depth)], q, a, [v]
+
+
+def t_niah_multikey_2(rng, depth):
+    """MK-NIAH-2: 8 distractor keys instead of 4."""
+    return t_niah_multikey(rng, depth, n_keys=8)
+
+
+def t_niah_multikey_3(rng, depth, n_keys=8):
+    """MK-NIAH-3: 8 keys AND UUID values."""
+    keys = rng.sample(WORDS, n_keys)
+    vals = [rnd_uuid(rng) for _ in keys]
+    tgt = rng.randrange(n_keys)
+    sents = [(f" The special magic uuid for {k} is {v}. ",
+              depth if i == tgt else rng.random())
+             for i, (k, v) in enumerate(zip(keys, vals))]
+    k, v = keys[tgt], vals[tgt]
+    q = f"\nQuestion: What is the special magic uuid for {k}?\n"
+    a = f"Answer: The special magic uuid for {k} is"
+    return sents, q, a, [v]
+
+
+def t_niah_multiquery(rng, depth, n_keys=4):
+    """MQ-NIAH: several keys planted, ALL of them queried at once. Distinct from
+    multikey (one target among distractors) and multivalue (one key, many
+    values): here every planted pair must be retrieved."""
+    keys = rng.sample(WORDS, n_keys)
+    vals = [rnd_number(rng) for _ in keys]
+    sents = [(f" The special magic number for {k} is {v}. ",
+              depth if i == 0 else rng.random())
+             for i, (k, v) in enumerate(zip(keys, vals))]
+    q = ("\nQuestion: What are the special magic numbers for "
+         + ", ".join(keys) + "?\n")
+    a = "Answer: The special magic numbers are"
+    return sents, q, a, vals          # every value must appear
+
+
+# ---- aggregation ----------------------------------------------------------
+# These return (body_text, question, answer_prefix, expected). build_one
+# tokenises body_text directly and does NOT use the haystack.
+
+CWE_COMMON, CWE_TOPK = 10, 10
+FWE_TOPK = 3
+
+
+def t_cwe(rng, depth, body_budget_tok=4096):
+    """CWE: a word list in which CWE_COMMON words recur many times and the rest
+    appear once. The model must name the top-K most common.
+
+    Word pool is built from WORDS x a digit suffix so every item is a single
+    common English token plus a digit -- keeps tokenisation tight and avoids
+    the model recognising real-word frequency from pretraining."""
+    pool = [f"{w}{i}" for w in WORDS for i in range(10)]
+    rng.shuffle(pool)
+    common = pool[:CWE_COMMON]
+    rest = pool[CWE_COMMON:]
+    # ~4 tokens per "word," item; reserve a little headroom.
+    n_items = max(CWE_COMMON * 6, int(body_budget_tok / 4.2))
+    reps = max(3, (n_items // 2) // max(1, CWE_COMMON))
+    items = [c for c in common for _ in range(reps)]
+    singles = rest * (max(0, n_items - len(items)) // max(1, len(rest)) + 1)
+    items += singles[: max(0, n_items - len(items))]
+    rng.shuffle(items)
+    body = " " + ", ".join(items) + ". "
+    q = f"\nQuestion: What are the {CWE_TOPK} most common words in the list above?\n"
+    a = "Answer: The most common words are"
+    return body, q, a, common          # all K must appear
+
+
+def t_fwe(rng, depth, body_budget_tok=4096, alpha=2.0):
+    """FWE: words drawn from a Zeta(alpha) distribution over a shuffled pool;
+    the model must name the top-3 most frequent. Unlike CWE the frequencies are
+    graded rather than two-tier, so ranking matters."""
+    pool = [f"{w}{i}" for w in WORDS for i in range(10)]
+    rng.shuffle(pool)
+    n_items = max(60, int(body_budget_tok / 4.2))
+    # Zeta: P(rank r) ~ r^-alpha over the shuffled pool.
+    wts = [(r + 1) ** (-alpha) for r in range(len(pool))]
+    tot = sum(wts)
+    cum, acc = [], 0.0
+    for w in wts:
+        acc += w / tot
+        cum.append(acc)
+    def draw():
+        u = rng.random()
+        lo, hi = 0, len(cum) - 1
+        while lo < hi:
+            mid = (lo + hi) // 2
+            if cum[mid] < u: lo = mid + 1
+            else: hi = mid
+        return pool[lo]
+    items = [draw() for _ in range(n_items)]
+    counts = {}
+    for it in items: counts[it] = counts.get(it, 0) + 1
+    top = [w for w, _ in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))][:FWE_TOPK]
+    rng.shuffle(items)
+    body = " " + ", ".join(items) + ". "
+    q = f"\nQuestion: What are the {FWE_TOPK} most frequently appearing words in the list above?\n"
+    a = "Answer: The most frequent words are"
+    return body, q, a, top
+
+
+
+# ---- QA with distractors (RULER qa_1 / qa_2) ------------------------------
+# RULER "extends traditional QA datasets by adding distracting paragraphs".
+# The GOLD paragraph sits at `depth`; the rest of the body is distractor
+# paragraphs drawn from other questions in the same dataset.
+#
+# HotpotQA yes/no answers are DROPPED: scoring is substring match on normalised
+# text, and "yes" matches almost any generation, which would score noise as a hit.
+
+QA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "qa_data")
+_QA_CACHE = {}
+
+
+def _load_qa(kind):
+    if kind in _QA_CACHE:
+        return _QA_CACHE[kind]
+    import pyarrow.parquet as pq
+    if kind == "squad":
+        t = pq.read_table(os.path.join(QA_DIR, "squad_val.parquet")).to_pydict()
+        rows = []
+        for ctx, qn, ans in zip(t["context"], t["question"], t["answers"]):
+            a = ans["text"][0] if ans["text"] else ""
+            if len(a) >= 3:
+                rows.append((ctx, qn, a))
+    else:
+        t = pq.read_table(os.path.join(QA_DIR, "hotpot_val.parquet")).to_pydict()
+        rows = []
+        for ctx, qn, a in zip(t["context"], t["question"], t["answer"]):
+            if a.lower() in ("yes", "no") or len(a) < 3:
+                continue                      # see note above
+            paras = [" ".join(sents) for sents in ctx["sentences"]]
+            rows.append(("\n\n".join(paras), qn, a))
+    if not rows:
+        raise SystemExit(f"no usable {kind} rows in {QA_DIR}")
+    _QA_CACHE[kind] = rows
+    return rows
+
+
+def _t_qa(kind, rng, depth, body_budget_tok):
+    rows = _load_qa(kind)
+    gold_ctx, question, answer = rows[rng.randrange(len(rows))]
+    gold = " " + gold_ctx.strip() + " "
+    # distractors from OTHER questions, appended until the budget is full
+    need = max(0, body_budget_tok - len(ENC.encode(gold)))
+    pre, post, got = [], [], 0
+    want_pre = int(max(0.0, min(1.0, depth)) * need)
+    while got < need:
+        d = rows[rng.randrange(len(rows))][0]
+        if d == gold_ctx:
+            continue
+        piece = " " + d.strip() + " "
+        n = len(ENC.encode(piece))
+        (pre if sum(len(ENC.encode(x)) for x in pre) < want_pre else post).append(piece)
+        got += n
+    body = "".join(pre) + gold + "".join(post)
+    q = f"\nQuestion: {question}\n"
+    a = "Answer:"
+    return body, q, a, [answer]
+
+
+def t_qa_1(rng, depth, body_budget_tok=4096):
+    """SQuAD paragraph + SQuAD distractors."""
+    return _t_qa("squad", rng, depth, body_budget_tok)
+
+
+def t_qa_2(rng, depth, body_budget_tok=4096):
+    """HotpotQA (multi-hop, 2 gold paragraphs) + distractors."""
+    return _t_qa("hotpot", rng, depth, body_budget_tok)
+
 TASKS = {
     "niah_single":     t_niah_single,
     "niah_multikey":   t_niah_multikey,
@@ -191,6 +404,16 @@ TASKS = {
     "vt":              t_vt,
     "val_forward":     t_val_forward,
     "val_backward":    t_val_backward,
+    # --- RULER parity (2026-09-17) ---
+    "niah_single_1":     t_niah_single_1,
+    "niah_single_3":     t_niah_single_3,
+    "niah_multikey_2":   t_niah_multikey_2,
+    "niah_multikey_3":   t_niah_multikey_3,
+    "niah_multiquery":   t_niah_multiquery,
+    "cwe":               t_cwe,
+    "fwe":               t_fwe,
+    "qa_1":              t_qa_1,
+    "qa_2":              t_qa_2,
 }
 
 
@@ -202,13 +425,30 @@ def build_one(task, rng, prompt_len, hay, depth, hay_off=0):
     replicates rather than independent draws -- which inflates apparent
     precision because the effective sample size is far below N.
     """
-    sents, q, a, expected = TASKS[task](rng, depth)
+    # BODY_TASKS (cwe/fwe) generate their own body and need to know the budget,
+    # so they are called with it; everything else keeps the original signature.
+    if task in BODY_TASKS:
+        probe_q, probe_a = TASKS[task](random.Random(0), depth, 64)[1:3]
+        tail0 = len(ENC.encode(probe_q)) + len(ENC.encode(probe_a))
+        body_text, q, a, expected = TASKS[task](rng, depth, max(64, prompt_len - tail0))
+    else:
+        sents, q, a, expected = TASKS[task](rng, depth)
     q_ids, a_ids = ENC.encode(q), ENC.encode(a)
     tail = len(q_ids) + len(a_ids)
     body_budget = prompt_len - tail
     if body_budget <= 0:
         raise SystemExit(f"--prompt-len {prompt_len} too small for task {task}")
 
+    if task in BODY_TASKS:
+        # The body IS the data. Repeat-or-truncate to hit the exact length; the
+        # C++ reader slices at a fixed stride so drift is fatal.
+        b = ENC.encode(body_text)
+        if len(b) < body_budget:
+            b = (b * (body_budget // max(1, len(b)) + 2))[:body_budget]
+        out = b[:body_budget] + q_ids + a_ids
+        assert len(out) == prompt_len, (len(out), prompt_len)
+        return out, expected
+
     # Plant each needle at its depth fraction of the BODY, then fill with haystack.
     planted = sorted(((max(0.0, min(1.0, d)), ENC.encode(s)) for s, d in sents),
                      key=lambda x: x[0])
@@ -255,8 +495,11 @@ def main():
 
     for task in tasks:
         rng = random.Random(args.seed)          # per-task, so tasks are independent
-        # (n + 2) windows so every prompt gets its own filler offset
-        hay = load_haystack(args.haystack, args.prompt_len * (args.n + 2))
+        # (n + 2) windows so every prompt gets its own filler offset.
+        # NOISE_TASKS (S-NIAH-1) deliberately use the repeated-noise filler
+        # instead of real text -- that IS the variant's definition in RULER.
+        hay_path = "" if task in NOISE_TASKS else args.haystack
+        hay = load_haystack(hay_path, args.prompt_len * (args.n + 2))
         bin_p = os.path.join(args.out_dir, f"{task}_prompts.bin")
         meta = []
         with open(bin_p, "wb") as f: