sieve_bench/common/data.py

"""SIEVE dataset builders. Fixed seeds; generated data shipped to data/ as json.

- load_corpus(n, seed): plain sentences from the on-box v2/SENT corpus
  (HF repo nickypro/sonar-sae, data_v2/sent_texts_*.json). Falls back to a
  shipped json snapshot when HF / network is unavailable.
- entity_presence_pairs(seed, n): cross-paraphrase entity-presence pairs.
- active_passive_set(seed, n): agent/patient active+passive constructions.

Every builder writes its output to data/<name>.json and reads it back on
subsequent calls so runs are reproducible without the corpus.
"""
from __future__ import annotations

import json
import os
import random

DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data")
HF_REPO = "nickypro/sonar-sae"


def _data_path(name: str) -> str:
    os.makedirs(DATA_DIR, exist_ok=True)
    return os.path.join(DATA_DIR, name)


# A bundled raw corpus snapshot shipped with the repo. This is the offline
# source of truth: it lets every corpus-backed task build WITHOUT HuggingFace,
# the on-box SONAR corpus, or any network access. A community user who only
# `pip install -r requirements.txt`'d gets a fully reproducible benchmark.
BUNDLED_CORPUS = "corpus_bundled.json"


# ---------------------------------------------------------------------------
# Corpus loader (on-box v2/SENT). Reuses tae_pool_ablate.load_sentences logic.
# Resolution order: (1) exact shipped/cached corpus_<n>_<seed>.json; (2) the
# on-box HF corpus (regenerates the exact file); (3) deterministic subsample of
# the BUNDLED snapshot shipped in data/ -> works fully offline, no SONAR/HF.
# ---------------------------------------------------------------------------
def load_corpus(n: int = 2000, min_words: int = 5, max_words: int = 60,
                seed: int = 0, max_shards: int = 6, cache_name: str = None):
    """Return list[str] of sentences. Caches to data/corpus_<n>_<seed>.json."""
    cache_name = cache_name or f"corpus_{n}_{seed}.json"
    fp = _data_path(cache_name)
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)

    # (2) on-box HF corpus, if reachable; (3) bundled-snapshot fallback.
    try:
        texts = _load_corpus_from_hf(n, min_words, max_words, seed, max_shards)
    except Exception:
        texts = _load_corpus_from_bundle(n, min_words, max_words, seed)

    with open(fp, "w") as f:
        json.dump(texts, f)
    return texts


def _load_corpus_from_bundle(n, min_words, max_words, seed):
    """Offline fallback: subsample the shipped bundled snapshot, honoring the
    word-length band so each task gets the distribution it expects. Deterministic
    in (n, seed). If fewer than n sentences pass the band filter we return all of
    them (so a run never crashes for lack of corpus)."""
    fp = _data_path(BUNDLED_CORPUS)
    if not os.path.exists(fp):
        raise RuntimeError(
            "Corpus unavailable: no HF access, no cached corpus_<n>_<seed>.json, "
            f"and the bundled snapshot {BUNDLED_CORPUS} is missing. Re-clone the "
            "repo (the snapshot ships in data/) or run on a box with HF access.")
    with open(fp) as f:
        pool = json.load(f)
    seen, cands = set(), []
    for s in pool:
        s = " ".join(str(s).split())
        if not (min_words <= len(s.split()) <= max_words):
            continue
        if s in seen:
            continue
        seen.add(s)
        cands.append(s)
    if len(cands) > n:
        cands = random.Random(seed).sample(cands, n)
    return cands


def _load_corpus_from_hf(n, min_words, max_words, seed, max_shards):
    from huggingface_hub import hf_hub_download
    cands, seen = [], set()
    for shard in range(max_shards):
        try:
            p = hf_hub_download(HF_REPO, f"data_v2/sent_texts_{shard:03d}.json",
                                repo_type="model")
        except Exception:
            break
        rows = json.load(open(p))
        for s in rows:
            s = " ".join(str(s).split())
            if not (min_words <= len(s.split()) <= max_words):
                continue
            if s in seen:
                continue
            seen.add(s)
            cands.append(s)
        if len(cands) >= 4 * n:
            break
    if not cands:
        raise RuntimeError(
            "Corpus unavailable (no HF access and no shipped snapshot). "
            "Run on the GPU box or ship data/corpus_*.json.")
    if len(cands) > n:
        cands = random.Random(seed).sample(cands, n)
    return cands


# ---------------------------------------------------------------------------
# A3 entity-presence paraphrase pairs (templated, deterministic).
# Each item: {text, entity, label(1=present), pid(paraphrase-group)}.
# Cross-paraphrase test: probe must read entity X regardless of surface form.
# ---------------------------------------------------------------------------
_ENTITIES = [
    "Amazon", "Google", "Tesla", "NASA", "Berlin", "Tokyo", "Brazil",
    "Microsoft", "Paris", "Canada", "Sony", "Egypt", "Apple", "Norway",
    "Reuters", "Oxford", "Boeing", "Kenya", "Toyota", "Sweden",
]

_TEMPLATES = [
    "The report mentioned that {E} announced a new initiative on Tuesday.",
    "According to officials, {E} will expand operations next year.",
    "Analysts said {E} had outperformed expectations this quarter.",
    "It was reported that {E} signed a major agreement abroad.",
    "Sources confirmed {E} plans to invest heavily in research.",
    "{E} was at the center of the discussion during the summit.",
]

# Distractor entities used to fill negative examples
_DISTRACTORS = ["Spain", "Intel", "Chile", "Ford", "London", "India"]


def entity_presence_pairs(seed: int = 0, n: int = 300):
    """Build entity-presence items for ONE fixed target entity, tested
    cross-paraphrase. label=1 -> the target entity is present (in some
    template surface form); label=0 -> a different entity fills that slot.

    Key design: the SAME target entity appears across many templates, and the
    cross-paraphrase split holds out TEMPLATES (surface forms), not entities.
    So a probe must read "target entity present" robustly to wording. The
    `pid` field is the template id (the paraphrase axis to hold out).
    Returns list of dicts: {text, entity(target), label, pid(template id)}.
    """
    fp = _data_path(f"entity_presence_{n}_{seed}.json")
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)

    rng = random.Random(seed)
    target = _ENTITIES[0]  # fixed target entity ("Amazon")
    negatives = [e for e in _ENTITIES[1:] + _DISTRACTORS if e != target]
    items = []
    # balance positives/negatives across templates; repeat with negative variety
    reps = max(1, n // (2 * len(_TEMPLATES)) + 1)
    for pid, tmpl in enumerate(_TEMPLATES):
        for _ in range(reps):
            items.append({"text": tmpl.format(E=target), "entity": target,
                          "label": 1, "pid": pid})
            neg = rng.choice(negatives)
            items.append({"text": tmpl.format(E=neg), "entity": target,
                          "label": 0, "pid": pid})
    rng.shuffle(items)
    items = items[:n]
    with open(fp, "w") as f:
        json.dump(items, f)
    return items


# ---------------------------------------------------------------------------
# A6 thematic-role: active/passive agent-patient. spaCy if available, else
# a deterministic templated generator (the default path on the CPU box).
# Each item: {text, construction(active|passive), agent, patient,
#             query_entity, is_agent(1/0)}.
# ---------------------------------------------------------------------------
_PEOPLE = [
    "Maria", "John", "Aisha", "Carlos", "Yuki", "Omar", "Elena", "David",
    "Priya", "Tom", "Sofia", "Hassan", "Anna", "Wei", "Lucas", "Fatima",
]
_VERBS = [
    ("chased", "was chased by"),
    ("helped", "was helped by"),
    ("called", "was called by"),
    ("praised", "was praised by"),
    ("blamed", "was blamed by"),
    ("followed", "was followed by"),
    ("warned", "was warned by"),
    ("rescued", "was rescued by"),
]


def active_passive_set(seed: int = 0, n: int = 300):
    """Build agent/patient pairs in both active and passive voice.

    For each (agent A, patient P, verb v):
      active:  "A v_active P."          agent=A patient=P
      passive: "P v_passive A."         agent=A patient=P (same semantics)
    We emit, per sentence, two query rows: query=agent (is_agent=1) and
    query=patient (is_agent=0). The benchmark trains on active, tests on
    passive: a surface-position probe nails active but fails passive.
    """
    fp = _data_path(f"active_passive_{n}_{seed}.json")
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)

    items = _build_active_passive(seed, n)
    with open(fp, "w") as f:
        json.dump(items, f)
    return items


# ---------------------------------------------------------------------------
# A6 NATURAL thematic-role (night-7 C2): real subject/object pairs mined from
# the v2 corpus with a spaCy full dependency parse, across diverse syntax
# (active / passive / relcl / ditrans / fronted). Each sentence carries ONE
# (agent, patient, verb, construction). The task emits per-entity QUERY rows
# (query=agent -> is_agent=1 ; query=patient -> is_agent=0) and trains on ONE
# construction (active), tests on a DIFFERENT construction -> cross-construction
# transfer is the binding test. Built offline from data/natural_relational_role.json
# (shipped); regenerated on-box with spaCy if absent.
# ---------------------------------------------------------------------------
NATURAL_ROLE_FILE = "natural_relational_role.json"


def natural_role_set(seed: int = 0, n: int = None):
    """Return list of QUERY-row dicts for the natural relational role task.

    Each underlying sentence yields TWO rows (one per entity), so a probe that
    is handed (z_sentence, entity_pointer) must decide whether the POINTED-TO
    entity is the agent. Rows:
      {text, construction, agent, patient, query_entity, is_agent(1/0), sid}
    Constructions: active / passive / relcl / ditrans / fronted (natural syntax).
    """
    fp = _data_path(NATURAL_ROLE_FILE)
    if not os.path.exists(fp):
        raise RuntimeError(
            f"natural role dataset {NATURAL_ROLE_FILE} missing; regenerate on a "
            "box with spaCy + HF corpus (see gen_natural_relational).")
    with open(fp) as f:
        blob = json.load(f)
    sents = blob["items"] if isinstance(blob, dict) else blob
    items = []
    for sid, s in enumerate(sents):
        ag, pa = s["agent"], s["patient"]
        constr = s["constr"]
        for q, lab in [(ag, 1), (pa, 0)]:
            items.append({"text": s["sent"], "construction": constr,
                          "agent": ag, "patient": pa, "query_entity": q,
                          "is_agent": lab, "verb": s.get("verb", ""), "sid": sid})
    if n is not None and len(items) > n:
        # keep whole sentences (pairs) together; truncate by sentence
        keep_sids = set(range(n // 2))
        items = [it for it in items if it["sid"] in keep_sids]
    return items


# ---------------------------------------------------------------------------
# A2 number-exact: sentences carrying one explicit numeric value. The probe
# (ridge) must recover the magnitude from z; exact-match on the rounded value.
# Each item: {text, value(float), pid(template id)}.
# ---------------------------------------------------------------------------
_NUM_TEMPLATES = [
    "The company reported {V} million dollars in revenue this quarter.",
    "Officials confirmed that {V} people attended the event downtown.",
    "The bridge spans {V} meters across the river valley.",
    "Researchers measured a temperature of {V} degrees during the test.",
    "The team scored {V} points in the final match of the season.",
    "About {V} students enrolled in the program this year.",
    "The shipment weighed {V} kilograms when it arrived at the port.",
    "Investors poured {V} thousand into the new startup last month.",
]


def number_exact_set(seed: int = 0, n: int = 300):
    """Build sentences with a single explicit integer value, varied across
    templates. label=the integer. Cross-template held-out so the probe reads
    the NUMBER not a template-specific artifact. Returns list of dicts:
    {text, value, pid}."""
    fp = _data_path(f"number_exact_{n}_{seed}.json")
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)
    rng = random.Random(seed)
    items = []
    # values spread over a wide magnitude range so r2 is meaningful and
    # corpus-median is an honest baseline.
    while len(items) < n:
        pid = rng.randrange(len(_NUM_TEMPLATES))
        v = rng.randint(1, 999)
        items.append({"text": _NUM_TEMPLATES[pid].format(V=v),
                      "value": int(v), "pid": pid})
    rng.shuffle(items)
    items = items[:n]
    with open(fp, "w") as f:
        json.dump(items, f)
    return items


# ---------------------------------------------------------------------------
# A4 negation-scope: a base clause that is either affirmed or negated, plus a
# distractor clause. The label is (1) negation present AND (2) it scopes the
# TARGET clause (not the distractor). Tests reading negation + its scope.
# Each item: {text, neg_present(0/1), neg_scopes_target(0/1), label, pid}.
# ---------------------------------------------------------------------------
_NEG_SUBJECTS = ["the manager", "the doctor", "the pilot", "the teacher",
                 "the engineer", "the lawyer", "the nurse", "the chef"]
_NEG_PREDS = [("approved the plan", "did not approve the plan"),
              ("signed the contract", "did not sign the contract"),
              ("attended the meeting", "did not attend the meeting"),
              ("finished the report", "did not finish the report"),
              ("accepted the offer", "did not accept the offer"),
              ("passed the exam", "did not pass the exam")]
_NEG_DISTRACT = ["while the office stayed open", "as the day went on",
                 "before the deadline arrived", "after the call ended",
                 "during the long afternoon", "once the rain stopped"]


def negation_scope_set(seed: int = 0, n: int = 300):
    """Build clauses that are affirmed/negated. Positive label = the TARGET
    clause is negated. To make the surface 'not'-detector NON-trivial as a
    *scope* test, half the negatives put 'not/no' into the distractor clause
    (negation present but NOT scoping the target) so a pure not-detector is
    fooled. Returns list of dicts: {text, neg_present, neg_scopes_target,
    label, pid}."""
    fp = _data_path(f"negation_scope_{n}_{seed}.json")
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)
    rng = random.Random(seed)
    items = []
    distract_neg = {"while the office stayed open": "while the office was not open",
                    "as the day went on": "as nothing went on",
                    "before the deadline arrived": "before no deadline arrived",
                    "after the call ended": "after the call had not ended",
                    "during the long afternoon": "during no long afternoon",
                    "once the rain stopped": "once the rain had not stopped"}
    while len(items) < n:
        subj = rng.choice(_NEG_SUBJECTS)
        pid = rng.randrange(len(_NEG_PREDS))
        aff, neg = _NEG_PREDS[pid]
        dist = rng.choice(_NEG_DISTRACT)
        kind = rng.randrange(3)
        if kind == 0:  # target negated -> label 1
            clause = neg
            d = dist
            label, np_, scope = 1, 1, 1
        elif kind == 1:  # target affirmed, no negation anywhere -> label 0
            clause = aff
            d = dist
            label, np_, scope = 0, 0, 0
        else:  # target affirmed, distractor negated -> negation present but NOT scoping target -> label 0
            clause = aff
            d = distract_neg[dist]
            label, np_, scope = 0, 1, 0
        text = f"{subj.capitalize()} {clause}, {d}."
        items.append({"text": text, "neg_present": np_,
                      "neg_scopes_target": scope, "label": label, "pid": pid})
    rng.shuffle(items)
    items = items[:n]
    with open(fp, "w") as f:
        json.dump(items, f)
    return items


# ---------------------------------------------------------------------------
# A9 coreference: a two-sentence passage with two named candidate entities and
# a pronoun; the query asks "does PRONOUN refer to ENTITY X". We use a
# gender/number heuristic so the antecedent is unambiguous from a HUMAN view
# but requires relational binding (which name the pronoun co-indexes) — a
# bag/surface reader cannot resolve which of the two named entities is meant.
# Each item: {text, query_entity, label(1=pronoun refers to query), pid}.
# ---------------------------------------------------------------------------
_COREF_MALE = ["John", "David", "Carlos", "Omar", "Lucas", "Tom", "Hassan", "Wei"]
_COREF_FEMALE = ["Maria", "Aisha", "Elena", "Priya", "Sofia", "Anna", "Fatima", "Yuki"]
_COREF_ACTIONS = [
    ("went to the office early", "was very tired"),
    ("finished the project", "felt proud"),
    ("called the client", "needed more time"),
    ("read the long report", "took a short break"),
    ("led the meeting", "answered every question"),
    ("drove to the airport", "missed the flight"),
]


def coreference_set(seed: int = 0, n: int = 300):
    """Build two-clause passages: 'A and B did X; PRONOUN did Y.' One of A,B is
    male and the other female; the pronoun (he/she) disambiguates the referent
    by gender. The query asks 'does the pronoun refer to <entity>'. The label
    flips with the pronoun, NOT with surface position (we randomize which of
    A,B is male and which slot they occupy). So a surface-position / bag reader
    cannot solve it; it needs to bind pronoun->name. Returns list of dicts:
    {text, query_entity, label, pid, pronoun}."""
    fp = _data_path(f"coreference_{n}_{seed}.json")
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)
    rng = random.Random(seed)
    items = []
    while len(items) < n:
        m = rng.choice(_COREF_MALE)
        f_ = rng.choice(_COREF_FEMALE)
        pid = rng.randrange(len(_COREF_ACTIONS))
        x, y = _COREF_ACTIONS[pid]
        # randomize surface order of the two names so position is decoupled
        if rng.random() < 0.5:
            first, second = m, f_
        else:
            first, second = f_, m
        # randomly choose which referent the pronoun picks out -> pronoun gender
        if rng.random() < 0.5:
            referent, pron = m, "He"
        else:
            referent, pron = f_, "She"
        text = f"{first} and {second} {x}; {pron} {y}."
        # emit two query rows (one per name) -> balanced labels
        for cand in (first, second):
            items.append({"text": text, "query_entity": cand,
                          "label": 1 if cand == referent else 0,
                          "pid": pid, "pronoun": pron})
    rng.shuffle(items)
    items = items[:n]
    with open(fp, "w") as f:
        json.dump(items, f)
    return items


# ---------------------------------------------------------------------------
# A8 length-generalization: entity-presence items binned SHORT vs LONG. A
# content probe (target entity present) is trained on SHORT and tested on LONG.
# Each item: {text, label, length_bin('short'|'long'), n_words}.
# ---------------------------------------------------------------------------
_LEN_TARGET = "Tesla"
_LEN_OTHERS = ["Amazon", "Google", "NASA", "Sony", "Boeing", "Toyota",
               "Microsoft", "Apple", "Reuters", "Oxford"]
_LEN_FILLERS = [
    "after a long and detailed discussion among the senior staff",
    "despite repeated concerns raised by several independent analysts",
    "while markets remained volatile across most major global regions",
    "as the quarterly figures were finally released to the public",
    "following months of careful negotiation behind closed doors",
    "amid growing speculation from reporters covering the industry",
]


def length_generalization_set(seed: int = 0, n: int = 300):
    """Entity-presence items where SHORT sentences are ~6-9 words and LONG ones
    ~18-24 words (extra clauses padded). The target entity sits at a RANDOM
    content slot in both bins. Probe trains short, tests long. label=target
    present. Returns list of dicts: {text, label, length_bin, n_words}."""
    fp = _data_path(f"length_gen_{n}_{seed}.json")
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)
    rng = random.Random(seed)
    items = []
    short_core = ["{E} announced new results today.",
                  "Reports said {E} expanded abroad.",
                  "Officials praised {E} this week.",
                  "{E} opened a new facility."]
    while len(items) < n:
        bin_ = "short" if (len(items) % 2 == 0) else "long"
        present = rng.random() < 0.5
        ent = _LEN_TARGET if present else rng.choice(_LEN_OTHERS)
        core = rng.choice(short_core).format(E=ent)
        if bin_ == "short":
            text = core
        else:
            f1 = rng.choice(_LEN_FILLERS)
            f2 = rng.choice([x for x in _LEN_FILLERS if x != f1])
            text = f"{core[:-1]}, {f1}, {f2}."
        items.append({"text": text, "label": 1 if present else 0,
                      "length_bin": bin_, "n_words": len(text.split())})
    rng.shuffle(items)
    items = items[:n]
    with open(fp, "w") as f:
        json.dump(items, f)
    return items


# ---------------------------------------------------------------------------
# Readable-property extractor (ported from tae_r9_meaning_coverage.readable_props).
# Used by t07 meaning-coverage. Returns a fixed-key dict of scalar floats.
# ---------------------------------------------------------------------------
import re as _re

_R9_POS = set("good great excellent wonderful happy joy love loved beautiful best better positive success win won amazing fantastic perfect nice enjoy glad pleased proud hope bright warm kind peace safe healthy strong rich free fun fine clear smart wise true honest".split())
_R9_NEG = set("bad terrible awful horrible miserable disgusting sad angry hate hated fear afraid worst worse negative fail failed failure loss lost pain hurt cruel ugly poor weak sick ill dead death kill killed war violence danger dangerous wrong false dark cold dirty broken difficult hard struggle suffer crisis problem disaster threat guilty corrupt".split())
_R9_NEGATORS = set("not no never none nobody nothing nowhere neither nor cannot can't won't don't doesn't didn't isn't aren't wasn't weren't hasn't haven't without".split())
_R9_P1 = set("i me my mine we us our ours myself ourselves".split())
_R9_P2 = set("you your yours yourself yourselves".split())
_R9_P3 = set("he him his she her hers it its they them their theirs himself herself itself themselves".split())
_R9_FUT = set("will shall going gonna".split())
_R9_PRES = set("is are am do does has have being".split())
_R9_PAST = set("was were did had been".split())
_R9_NUMW = set("one two three four five six seven eight nine ten eleven twelve hundred thousand million billion".split())
_R9_STARTERS = set("the a an this that these those it he she they we you i but and or so if when while".split())
_R9_tokre = _re.compile(r"[A-Za-z']+|[0-9]+|[.,!?;:]")
_R9_wordre = _re.compile(r"[A-Za-z']+")


def readable_props(text):
    """Scalar readable-property features (sentiment/tense/number/negation/
    coref/NE/discourse/length). Ported from tae_r9_meaning_coverage."""
    toks = _R9_tokre.findall(text)
    ws = [t.lower() for t in toks if _R9_wordre.fullmatch(t)]
    n = max(len(ws), 1)
    f = {}
    pos = sum(w in _R9_POS for w in ws); neg = sum(w in _R9_NEG for w in ws)
    f["sent_pos"] = pos / n; f["sent_neg"] = neg / n
    f["sent_signed"] = (pos - neg) / n; f["sent_mag"] = (pos + neg) / n
    f["tense_ed"] = sum(bool(_re.fullmatch(r"[a-z]+ed", w)) for w in ws) / n
    f["tense_ing"] = sum(bool(_re.fullmatch(r"[a-z]+ing", w)) for w in ws) / n
    f["tense_aux_past"] = sum(w in _R9_PAST for w in ws) / n
    f["tense_aux_pres"] = sum(w in _R9_PRES for w in ws) / n
    f["tense_fut"] = sum(w in _R9_FUT for w in ws) / n
    f["num_plural"] = sum(bool(_re.fullmatch(r"[a-z]{3,}s", w)) and not w.endswith("ss") for w in ws) / n
    f["num_words"] = sum(w in _R9_NUMW for w in ws) / n
    f["num_digits"] = sum(bool(_re.fullmatch(r"[0-9]+", t)) for t in toks) / n
    neg_idx = [i for i, w in enumerate(ws) if w in _R9_NEGATORS]
    f["neg_present"] = 1.0 if neg_idx else 0.0
    f["neg_count"] = len(neg_idx) / n
    f["neg_first_pos"] = (neg_idx[0] / n) if neg_idx else -1.0
    f["pron_1"] = sum(w in _R9_P1 for w in ws) / n
    f["pron_2"] = sum(w in _R9_P2 for w in ws) / n
    f["pron_3"] = sum(w in _R9_P3 for w in ws) / n
    f["pron_any"] = 1.0 if any(w in _R9_P1 | _R9_P2 | _R9_P3 for w in ws) else 0.0
    caps_mid = sum(1 for i, t in enumerate(toks)
                   if _R9_wordre.fullmatch(t) and t[0].isupper() and i > 0
                   and t.lower() not in _R9_STARTERS)
    f["ne_caps"] = caps_mid / n
    f["len_tok"] = len(ws) / 30.0
    f["punc_comma"] = sum(t == "," for t in toks) / n
    f["punc_q"] = 1.0 if "?" in toks else 0.0
    f["punc_excl"] = 1.0 if "!" in toks else 0.0
    return f


def _build_active_passive(seed, n):
    """Role-balanced generator: for every (agent=A,patient=P,verb) we ALSO emit
    the swapped (agent=P,patient=A) so each name is agent exactly as often as
    patient. This removes any name-frequency confound from the query-entity
    feature -> the query one-hot carries ZERO marginal signal, so only true
    thematic binding (absent in a bag) could lift the score. Without this the
    A6 'bag~chance' guarantee leaks ~0.67 via name priors."""
    rng = random.Random(seed)
    items = []
    sid = 0
    while len(items) < n:
        a, p = rng.sample(_PEOPLE, 2)
        v_act, v_pas = rng.choice(_VERBS)
        for ag, pa in [(a, p), (p, a)]:  # both role assignments -> balanced
            active = f"{ag} {v_act} {pa}."
            passive = f"{pa} {v_pas} {ag}."
            for text, constr in [(active, "active"), (passive, "passive")]:
                items.append({"text": text, "construction": constr,
                              "agent": ag, "patient": pa,
                              "query_entity": ag, "is_agent": 1, "sid": sid})
                items.append({"text": text, "construction": constr,
                              "agent": ag, "patient": pa,
                              "query_entity": pa, "is_agent": 0, "sid": sid})
            sid += 1
    rng.shuffle(items)
    return items[:n]


# ---------------------------------------------------------------------------
# GENERATIVE-TRACK builders (C / D families). All templated + deterministic so
# they run without the corpus; corpus-backed variants fall back gracefully.
# ---------------------------------------------------------------------------

# A small stop-word set used to pull "content words" out of a sentence for the
# sentence-from-words (C1) task and the lexical-bag scoring of edits.
_STOP = {
    "the", "a", "an", "of", "to", "in", "on", "at", "and", "or", "but", "for",
    "with", "as", "by", "that", "this", "it", "its", "is", "was", "were", "are",
    "be", "been", "had", "has", "have", "will", "would", "can", "could", "from",
    "their", "his", "her", "they", "he", "she", "we", "you", "i", "into", "out",
    "up", "down", "over", "than", "then", "so", "if", "not", "no", "do", "did",
    "during", "next", "new", "said", "according", "reported", "confirmed",
}


def content_words(text: str):
    """Lowercased content tokens (stopwords + pure punctuation removed)."""
    out = []
    for w in text.split():
        c = w.strip(".,!?;:'\"()").lower()
        if c and c not in _STOP and not c.isdigit():
            out.append(c)
    return out


def sentence_from_words_items(seed: int = 0, n: int = 300, smoke: bool = False):
    """C1: take real corpus sentences, expose their SHUFFLED content words.
    item: {gold(sentence), words(shuffled content words)}."""
    fp = _data_path(f"words_from_sentence_{n}_{seed}.json")
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)
    sents = load_corpus(n=max(n * 3, 600), seed=seed, min_words=6, max_words=22)
    rng = random.Random(seed)
    items = []
    for s in sents:
        w = content_words(s)
        if not (3 <= len(w) <= 12):
            continue
        ws = w[:]
        rng.shuffle(ws)
        items.append({"gold": s, "words": ws})
        if len(items) >= n:
            break
    with open(fp, "w") as f:
        json.dump(items, f)
    return items


# Frequency-stratified vocabulary for the round-trip task (C2). Words are tagged
# with a coarse frequency band so accuracy can be reported vs frequency.
_VOCAB_BANDS = {
    "high": ["time", "people", "world", "year", "government", "country", "city",
             "company", "water", "money", "school", "family", "music", "story",
             "market", "health", "system", "power", "energy", "history"],
    "mid": ["glacier", "senator", "vaccine", "harvest", "orchestra", "comet",
            "monastery", "turbine", "sediment", "ballot", "pension", "refinery",
            "antibody", "tournament", "estuary", "manuscript", "satellite",
            "diplomat", "fossil", "canyon"],
    "low": ["quokka", "obsidian", "samovar", "zephyr", "philately", "tessellation",
            "quasar", "marzipan", "cassowary", "gabbro", "petrichor", "syzygy",
            "thylacine", "cuneiform", "axolotl", "mycelium", "narwhal",
            "basalt", "lemur", "lichen"],
}


def vocab_set(seed: int = 0, n: int = None, smoke: bool = False):
    """C2: frequency-stratified single words. item: {word, band}."""
    fp = _data_path("vocab_bands.json")
    if os.path.exists(fp):
        with open(fp) as f:
            items = json.load(f)
    else:
        items = []
        for band, words in _VOCAB_BANDS.items():
            for w in words:
                items.append({"word": w, "band": band})
        with open(fp, "w") as f:
            json.dump(items, f)
    if n:
        items = items[:n]
    return items


# Concept lexicons for steering / injection / role tasks. A sentence is
# "positive" for a concept if it contains one of its trigger words.
CONCEPT_LEX = {
    "money": ["money", "dollars", "price", "cost", "bank", "profit", "loan",
              "tax", "salary", "budget", "investment", "cash", "payment"],
    "location": ["city", "country", "mountain", "river", "ocean", "village",
                 "street", "border", "island", "desert", "forest", "harbor"],
    "time": ["yesterday", "tomorrow", "century", "morning", "evening", "decade",
             "hour", "minute", "week", "month", "summer", "winter", "midnight"],
    "weather": ["rain", "storm", "snow", "sunny", "wind", "cloud", "fog",
                "thunder", "drought", "humid", "frost", "hail"],
}


def _concept_label(text, concept):
    toks = set(content_words(text))
    return int(any(w in toks for w in CONCEPT_LEX[concept]))


def concept_corpus(seed: int = 0, n: int = 600, concept: str = "money",
                   smoke: bool = False):
    """D4/D5: corpus sentences labeled present/absent for `concept`.
    item: {text, label}. Balanced positive/negative where possible."""
    fp = _data_path(f"concept_{concept}_{n}_{seed}.json")
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)
    sents = load_corpus(n=max(n * 5, 2000), seed=seed, min_words=6, max_words=30)
    pos = [s for s in sents if _concept_label(s, concept)]
    neg = [s for s in sents if not _concept_label(s, concept)]
    rng = random.Random(seed)
    rng.shuffle(pos)
    rng.shuffle(neg)
    half = n // 2
    pos, neg = pos[:half], neg[:half]
    items = [{"text": s, "label": 1} for s in pos] + \
            [{"text": s, "label": 0} for s in neg]
    rng.shuffle(items)
    with open(fp, "w") as f:
        json.dump(items, f)
    return items


# Multi-sentence builders (D2 edit-2-of-3, D3 reorder, D6 role swap). Built from
# short, self-contained templated clauses so per-sentence semantics are clean.
_SUBJECTS = ["The scientist", "A teacher", "The pilot", "A farmer", "The nurse",
             "A lawyer", "The chef", "An artist", "The engineer", "A sailor"]
_PREDICATES = [
    "studied the ancient ruins", "repaired the old engine",
    "planted rows of wheat", "painted a bright mural",
    "measured the rising tide", "wrote a long letter",
    "fixed the broken radio", "watered the green garden",
    "counted the silver coins", "climbed the steep hill",
    "read the thick report", "cooked a warm meal",
]


def _clause(rng):
    return f"{rng.choice(_SUBJECTS)} {rng.choice(_PREDICATES)}"


def three_sentence_items(seed: int = 0, n: int = 200, smoke: bool = False):
    """D2/D3: 3-sentence docs with a clean swap/edit target.
    item: {sents:[s1,s2,s3], text, s2_alt(replacement clause for s2)}."""
    fp = _data_path(f"three_sentence_{n}_{seed}.json")
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)
    rng = random.Random(seed)
    items = []
    for _ in range(n):
        s = [_clause(rng) + "." for _ in range(3)]
        # s2_alt: a distinct replacement clause for sentence 2 only
        alt = _clause(rng) + "."
        while alt == s[1]:
            alt = _clause(rng) + "."
        items.append({"sents": s, "text": " ".join(s), "s2_alt": alt})
    with open(fp, "w") as f:
        json.dump(items, f)
    return items


# ---------------------------------------------------------------------------
# RECONSTRUCTED builders for the parallel CORE-track tasks (t07/t08/t09). These
# match the field contracts those tasks consume; kept here so a single data.py
# serves both tracks.
# ---------------------------------------------------------------------------

# A8 length-generalization (v0.3 natural-corpus binding-beyond-bag rework).
#
# HISTORY. v0.1 used one LITERAL token ("Amazon") -> a bag read it at any length
# so score==bag==1.0 (vacuous). v0.2 used a money-vs-nature CONCEPT with disjoint
# short/long trigger vocab in 4 fixed TEMPLATES; the bag floor dropped to ~0 (good)
# but the concept was so cleanly separable in EVERY encoder that the short->long z
# AUC hit the within-long ceiling for all 9 encoders -> normalized score==1.0 for
# all (cross-encoder std .005), so the task still carried NO ranking signal and was
# excluded from the LIVE intersection.
#
# v0.3 keeps the binding-beyond-bag idea but makes the concept LIVE in REAL,
# lexically-varied corpus context so that LENGTH genuinely DILUTES the concept
# signal and the short-trained readout direction must GENERALIZE across length:
#   - target = a paraphrase-robust SEMANTIC concept (FINANCE/MONEY topic), labelled
#     by a money lexicon, with the POSITIVE trigger vocab DISJOINT short vs long
#     (short-train words never appear in the long-test bin) -> a bag CANNOT
#     short->long transfer (held-out words) so the bag floor is ~0.5;
#   - the concept clause is EMBEDDED inside REAL corpus sentences drawn from the
#     v2 corpus: SHORT items are a single short corpus sentence carrying (or not)
#     a short money word; LONG items prepend/append 2-3 unrelated long corpus
#     sentences so the money word is ONE content token among ~25-40 -> 1/N
#     dilution. An encoder that BINDS the concept (keeps it salient + in a
#     length-stable direction) transfers short->long; a length-fragile / pure
#     averaging encoder sees the direction rotate/dilute and DEGRADES.
#   - score = short-trained -> LONG z AUC, normalized into [surface-position null,
#     within-long ceiling]: it measures GENERALIZATION-across-length, not raw
#     readability. The within-long ceiling factors OUT how hard the diluted long
#     read is in absolute terms, so the headline is pure length-transfer.
#
# Each item: {text, label(1=money concept present), length_bin('short'|'long'),
#             trigger(the money/other word injected), n_words}.

# concept = MONEY/FINANCE. Trigger vocab split so short-train words != long-test
# -> a literal bag trained on SHORT money words cannot fire on LONG money words.
_LG_MONEY_SHORT = ["dollars", "salary", "budget", "loan", "profit", "tax",
                   "wages", "savings"]
_LG_MONEY_LONG = ["revenue", "investment", "dividend", "mortgage", "earnings",
                  "funding", "payroll", "subsidy", "capital", "pension"]
# negative (non-money) injected content words, also split short/long, so the
# NEGATIVE class is ALSO a held-out content word (not just "no money word"):
# this keeps the bag from solving negatives by the absence of a known word.
_LG_OTHER_SHORT = ["garden", "mountain", "river", "forest", "harbor", "meadow",
                   "orchard", "valley"]
_LG_OTHER_LONG = ["glacier", "canyon", "tundra", "lagoon", "savanna",
                  "plateau", "estuary", "prairie", "fjord", "marsh"]

# Carrier clauses (short, neutral) into which the trigger word is injected at a
# varied content slot. Real corpus sentences supply the LONG distractor context.
_LG_CARRIERS = [
    "the {W} was mentioned in the notes",
    "they finally discussed the {W} today",
    "a report described the {W} in brief",
    "everyone kept asking about the {W}",
    "the {W} came up during the review",
    "she wrote a short line about the {W}",
]


def length_generalization_set(seed: int = 0, n: int = 2000, smoke: bool = False):
    """v0.3: money-concept presence embedded in REAL corpus context, binned
    SHORT vs LONG, with DISJOINT short/long trigger vocab. Probe trains short,
    tests long; the score measures whether the concept readout GENERALIZES across
    length (binding-beyond-bag), normalized into the within-long ceiling."""
    fp = _data_path(f"length_gen_{n}_{seed}.json")
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)
    rng = random.Random(seed)
    # Real corpus distractor sentences for the LONG bin (lexically varied).
    try:
        pool = load_corpus(n=max(n, 1500), seed=seed, min_words=8, max_words=22)
    except Exception:
        pool = []
    if len(pool) < 50:
        # offline fallback: synthetic neutral distractors (still lexically varied)
        pool = [c.format(W=w) for c in _LG_CARRIERS
                for w in ("meeting", "weather", "schedule", "building", "journey",
                          "festival", "machine", "lecture", "harvest", "voyage")]
    rng.shuffle(pool)

    def _make(label, length_bin):
        money = (label == 1)
        if length_bin == "short":
            words = _LG_MONEY_SHORT if money else _LG_OTHER_SHORT
        else:
            words = _LG_MONEY_LONG if money else _LG_OTHER_LONG
        w = rng.choice(words)
        clause = rng.choice(_LG_CARRIERS).format(W=w)
        if length_bin == "short":
            text = clause[0].upper() + clause[1:] + "."
        else:
            # embed the concept clause among 2-3 real corpus distractor sentences
            k = rng.randint(2, 3)
            distractors = [pool[rng.randrange(len(pool))] for _ in range(k)]
            parts = distractors[:]
            insert = rng.randint(0, len(parts))  # varied position of the clause
            parts.insert(insert, clause[0].upper() + clause[1:])
            text = ". ".join(p.rstrip(".") for p in parts) + "."
        return {"text": text, "label": label, "length_bin": length_bin,
                "trigger": w, "n_words": len(text.split())}

    items = []
    while len(items) < n:
        for length_bin in ("short", "long"):
            items.append(_make(1, length_bin))
            items.append(_make(0, length_bin))
    rng.shuffle(items)
    items = items[:n]
    with open(fp, "w") as f:
        json.dump(items, f)
    return items


# A9 coreference: two-name passage where a gendered pronoun disambiguates the
# referent, and the surface ORDER of the two names is randomized so position
# carries no cue. label = 1 if the query_entity is the pronoun's referent.
_MALE = ["John", "Carlos", "Omar", "David", "Tom", "Hassan", "Lucas", "Wei"]
_FEMALE = ["Maria", "Aisha", "Yuki", "Elena", "Priya", "Sofia", "Anna", "Fatima"]
_COREF_TMPL = [
    "{n1} met {n2} at the station, and {pro} carried the heavy suitcase.",
    "{n1} sat near {n2} during lunch, then {pro} left the room quietly.",
    "{n1} called {n2} last night because {pro} needed urgent help.",
    "{n1} thanked {n2} warmly after {pro} finished the difficult task.",
]


def coreference_set(seed: int = 0, n: int = 2000, smoke: bool = False):
    fp = _data_path(f"coreference_{n}_{seed}.json")
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)
    rng = random.Random(seed)
    items = []
    while len(items) < n:
        male = rng.choice(_MALE)
        female = rng.choice(_FEMALE)
        # pronoun picks the referent by gender
        referent, pro = rng.choice([(male, "he"), (female, "she")])
        # randomize surface order of the two names -> position carries no cue
        n1, n2 = (male, female) if rng.random() < 0.5 else (female, male)
        tmpl = rng.choice(_COREF_TMPL)
        text = tmpl.format(n1=n1, n2=n2, pro=pro)
        # emit two query rows: referent (label 1) and the other name (label 0)
        other = female if referent == male else male
        items.append({"text": text, "query_entity": referent, "label": 1})
        items.append({"text": text, "query_entity": other, "label": 0})
    rng.shuffle(items)
    items = items[:n]
    with open(fp, "w") as f:
        json.dump(items, f)
    return items


# NOTE: the A7 readable_props extractor is the RICH R9 version defined earlier
# in this file (30 features, ported from tae_r9_meaning_coverage). A thin
# 10-feature duplicate was removed here during the merge so the R9 version is
# the one t07 (meaning-coverage) uses.


# Transitive SVO triples for the causal-identifiability role-swap (D6).
def svo_swap_items(seed: int = 0, n: int = 200, smoke: bool = False):
    """D6: single-sentence SVO docs with a role-swapped counterfactual.
    item: {text(A V B), swapped(B V A), agent, patient, verb}."""
    fp = _data_path(f"svo_swap_{n}_{seed}.json")
    if os.path.exists(fp):
        with open(fp) as f:
            return json.load(f)
    rng = random.Random(seed)
    verbs = ["helped", "chased", "called", "praised", "blamed", "followed",
             "warned", "rescued", "questioned", "interviewed"]
    items = []
    for _ in range(n):
        a, p = rng.sample(_PEOPLE, 2)
        v = rng.choice(verbs)
        items.append({"text": f"{a} {v} {p}.", "swapped": f"{p} {v} {a}.",
                      "agent": a, "patient": p, "verb": v})
    with open(fp, "w") as f:
        json.dump(items, f)
    return items


# ---------------------------------------------------------------------------
# D1 word-edit (v0.2 rework). The v0.1 task swapped ONE entity pair
# (Amazon->Google) over 6 templates -> a single trivially-editable named-entity
# direction in SONAR (diff-of-means == naive-add), so score==base==ceiling==1.0.
# The rework uses a LARGER, MORE DIVERSE set of word pairs spanning parts of
# speech (noun/verb/adjective), an abstract<->concrete axis, and rare/uncommon
# words, each instantiated in MULTIPLE neutral carrier sentences so the edit
# direction is learned over varied contexts. This makes target-success vs
# collateral-preservation a non-trivial trade-off that separates a held-out
# diff-of-means edit from a naive single-vector add.
# Each pair: {x, y, pos, axis}. Carriers are generic so the swapped word is the
# only content difference -> clean collateral measurement.
# ---------------------------------------------------------------------------
# Pairs tagged with the carrier-set appropriate to their part of speech, so the
# swapped word reads naturally (clean SONAR decode + clean collateral measure).
_WORD_EDIT_PAIRS = [
    # concrete nouns
    {"x": "river", "y": "mountain", "pos": "noun", "axis": "concrete"},
    {"x": "doctor", "y": "teacher", "pos": "noun", "axis": "concrete"},
    {"x": "engine", "y": "garden", "pos": "noun", "axis": "concrete"},
    {"x": "castle", "y": "harbor", "pos": "noun", "axis": "concrete"},
    {"x": "winter", "y": "summer", "pos": "noun", "axis": "concrete"},
    # abstract nouns
    {"x": "freedom", "y": "justice", "pos": "noun", "axis": "abstract"},
    {"x": "sorrow", "y": "courage", "pos": "noun", "axis": "abstract"},
    {"x": "wisdom", "y": "silence", "pos": "noun", "axis": "abstract"},
    # rare / uncommon nouns
    {"x": "glacier", "y": "savanna", "pos": "noun", "axis": "rare"},
    {"x": "obsidian", "y": "marble", "pos": "noun", "axis": "rare"},
    {"x": "monastery", "y": "laboratory", "pos": "noun", "axis": "rare"},
    {"x": "diplomat", "y": "merchant", "pos": "noun", "axis": "rare"},
    {"x": "turbine", "y": "lantern", "pos": "noun", "axis": "rare"},
    # verbs (past tense, fit a "they {W} ..." carrier)
    {"x": "increased", "y": "decreased", "pos": "verb", "axis": "antonym"},
    {"x": "praised", "y": "criticized", "pos": "verb", "axis": "antonym"},
    {"x": "arrived", "y": "departed", "pos": "verb", "axis": "antonym"},
    # adjectives (fit a "the ... was {W}" carrier)
    {"x": "ancient", "y": "modern", "pos": "adj", "axis": "antonym"},
    {"x": "bright", "y": "gloomy", "pos": "adj", "axis": "antonym"},
    {"x": "fragile", "y": "sturdy", "pos": "adj", "axis": "antonym"},
]

# POS-appropriate carriers. The swapped word is the only content difference.
_WORD_EDIT_CARRIERS = {
    "noun": [
        "The {W} was clearly visible from the road.",
        "Everyone in the town talked about the {W} that morning.",
        "She wrote a long essay describing the {W} in detail.",
        "According to the guide, the {W} drew many visitors.",
        "They spent the afternoon thinking about the {W}.",
        "A short article mentioned the {W} near the end.",
    ],
    "verb": [
        "The visitors {W} just before noon on Tuesday.",
        "Reports say the workers {W} during the long meeting.",
        "Last spring the travellers {W} without much warning.",
        "Officials noted that the numbers {W} that year.",
        "Quietly, the guests {W} as the evening went on.",
        "By midday the crowd {W} along the river.",
    ],
    "adj": [
        "The building looked remarkably {W} from the hill.",
        "Visitors found the old hall surprisingly {W} inside.",
        "Critics described the new design as rather {W}.",
        "The whole valley seemed {W} in the morning light.",
        "Reporters called the proposal unexpectedly {W}.",
        "Everyone agreed the room felt {W} that day.",
    ],
}


def word_edit_pairs(seed: int = 0, n: int = None, smoke: bool = False):
    """D1: diverse word-swap pairs (noun/verb/adj, concrete/abstract/rare) in
    POS-appropriate carriers. Returns dict with `pairs` (the X->Y definitions)
    and `items` (per-pair, per-carrier instantiations). item:
    {x_text, y_text, x, y, pos, axis, cid}."""
    pairs = _WORD_EDIT_PAIRS
    items = []
    for p in pairs:
        for cid, carrier in enumerate(_WORD_EDIT_CARRIERS[p["pos"]]):
            items.append({
                "x_text": carrier.format(W=p["x"]),
                "y_text": carrier.format(W=p["y"]),
                "x": p["x"], "y": p["y"], "pos": p["pos"], "axis": p["axis"],
                "cid": cid,
            })
    return {"pairs": pairs, "items": items}