data.py

"""Generate synthetic lessons and an eval suite of invented-word facts.

Design goals:
- Zero internet contamination: all words are randomly generated syllable strings.
- Composition chains across lessons: later lessons reuse earlier lessons'
  entities as subjects so 2-hop chains naturally span lesson boundaries.
- Balanced "negative" category: half the negatives are actually real facts
  (gold_is_true=True). This prevents an "always False" trivial strategy from
  scoring 50% and forces architectures to actually discriminate.
- Deterministic: fully reproducible from a seed.
- Round-trippable: `load_dataset_json` deserializes a previously serialized
  `dataset.json` back into a `Dataset`, so `runs.replay` can verify against
  the exact dataset that was recorded without re-running the generator.
"""

from __future__ import annotations
import json
import random
from pathlib import Path

from tasks import Dataset, Fact, Lesson, Question


_CONS = "bdfgklmnprstvz"
_VOW = "aeiou"


def _syllable(rng: random.Random) -> str:
    s = rng.choice(_CONS) + rng.choice(_VOW)
    if rng.random() < 0.35:
        s += rng.choice(_CONS)
    return s


def _gen_word(rng: random.Random) -> str:
    n = rng.randint(2, 3)
    return "".join(_syllable(rng) for _ in range(n))


def _gen_unique_words(rng: random.Random, n: int) -> list[str]:
    seen: set[str] = set()
    out: list[str] = []
    while len(out) < n:
        w = _gen_word(rng)
        if w in seen or len(w) < 4:
            continue
        seen.add(w)
        out.append(w)
    return out


def generate_dataset(
    seed: int = 42,
    n_lessons: int = 10,
    facts_per_lesson: int = 10,
    n_relations: int = 6,
    n_negatives: int = 150,
    n_compositions: int = 100,
    n_compositions_3hop: int = 100,
) -> Dataset:
    rng = random.Random(seed)

    entities_per_lesson = facts_per_lesson
    n_entities = n_lessons * entities_per_lesson + 10  # small reserve

    relation_words = _gen_unique_words(rng, n_relations)
    entity_words = _gen_unique_words(rng, n_entities)

    vocab: dict[int, str] = {}
    relation_ids: list[int] = []
    entity_ids: list[int] = []
    for w in relation_words:
        tid = len(vocab); vocab[tid] = w; relation_ids.append(tid)
    for w in entity_words:
        tid = len(vocab); vocab[tid] = w; entity_ids.append(tid)

    # --- build lessons as a growing knowledge graph ------------------------
    # Each lesson introduces new entities as objects, with subjects sampled
    # from the pool of previously-taught entities. This guarantees 2-hop
    # composition chains that span lesson boundaries.
    lessons: list[Lesson] = []
    all_facts_by_lesson: list[list[Fact]] = []
    known: list[int] = []
    fact_index: dict[tuple[int, int], tuple[int, int]] = {}  # (s,r) -> (o, lesson_idx)

    entity_cursor = 0
    for li in range(n_lessons):
        facts: list[Fact] = []
        attempts = 0
        while len(facts) < facts_per_lesson:
            attempts += 1
            if attempts > 1000:
                raise RuntimeError("unable to generate enough non-colliding facts")
            if entity_cursor >= len(entity_ids):
                raise RuntimeError("ran out of entities — increase reserve")
            new_obj = entity_ids[entity_cursor]
            subj_pool = known if known else [entity_ids[entity_cursor + 1]]
            s = rng.choice(subj_pool)
            r = rng.choice(relation_ids)
            if (s, r) in fact_index:
                # avoid collisions that would create ambiguous answers
                continue
            entity_cursor += 1
            f = Fact(subject=s, relation=r, object=new_obj)
            facts.append(f)
            fact_index[(s, r)] = (new_obj, li)
            known.append(new_obj)
        lessons.append(Lesson(idx=li, facts=facts))
        all_facts_by_lesson.append(facts)

    # --- questions ---------------------------------------------------------
    questions: list[Question] = []
    qid = 0

    # 1. Recall: one question per fact.
    for li, facts in enumerate(all_facts_by_lesson):
        for f in facts:
            questions.append(Question(
                qid=qid, kind="recall",
                subject=f.subject, relation=f.relation,
                source_lessons=(li,),
                gold_object=f.object,
            ))
            qid += 1

    # 2. Composition: (s, r1, r2) -> o chains spanning different lessons.
    by_subject: dict[int, list[tuple[Fact, int]]] = {}
    for li, facts in enumerate(all_facts_by_lesson):
        for f in facts:
            by_subject.setdefault(f.subject, []).append((f, li))
    chains: list[tuple[Fact, int, Fact, int]] = []
    for li1, facts in enumerate(all_facts_by_lesson):
        for f1 in facts:
            for (f2, li2) in by_subject.get(f1.object, []):
                if li1 == li2:
                    continue
                chains.append((f1, li1, f2, li2))
    rng.shuffle(chains)
    for (f1, li1, f2, li2) in chains[:n_compositions]:
        questions.append(Question(
            qid=qid, kind="composition",
            subject=f1.subject, relation=f1.relation, relation2=f2.relation,
            source_lessons=tuple(sorted({li1, li2})),
            gold_object=f2.object,
        ))
        qid += 1

    # 3. Negatives: balanced fact verification.
    #    Half are real facts (gold_is_true=True), half are wrong (gold_is_true=False).
    #    The two halves are drawn from independent shuffles of the same
    #    fact pool, so a fact can serve as the basis for both a True and
    #    a False question (different questions, different qids). An
    #    architecture that always answers "False" scores 50% — it has to
    #    actually discriminate.
    #
    #    History: a prior version of this code used `verification[:half]`
    #    then `verification[half:n_negatives]` from a single shuffle,
    #    which silently produced 75/25 imbalanced negatives when
    #    `n_negatives=150` and the verification pool had only 100 facts.
    #    That bug was fixed in cycle 6 (see paper/paper.md §5.6).
    verification: list[tuple[Fact, int]] = []
    for li, facts in enumerate(all_facts_by_lesson):
        for f in facts:
            verification.append((f, li))

    half = n_negatives // 2
    if half > len(verification):
        raise ValueError(
            f"n_negatives // 2 = {half} exceeds verification pool size "
            f"{len(verification)}. Either lower n_negatives or enable "
            f"sampling with replacement in your own fork."
        )

    # Independent shuffle for True half
    rng.shuffle(verification)
    for (f, li) in verification[:half]:
        questions.append(Question(
            qid=qid, kind="negative",
            subject=f.subject, relation=f.relation,
            candidate=f.object,
            source_lessons=(li,),
            gold_is_true=True,
        ))
        qid += 1

    # Independent shuffle for False half (may overlap with True half)
    rng.shuffle(verification)
    for (f, li) in verification[:half]:
        wrong_pool = [e for e in entity_ids if e != f.object]
        wrong = rng.choice(wrong_pool)
        questions.append(Question(
            qid=qid, kind="negative",
            subject=f.subject, relation=f.relation,
            candidate=wrong,
            source_lessons=(li,),
            gold_is_true=False,
        ))
        qid += 1

    # 4. 3-hop composition: (s, r1, r2, r3) -> o3 across three distinct lessons.
    # Added in cycle 20 (see paper §3.9). A 3-hop chain (f1, f2, f3) requires
    # f1.object == f2.subject, f2.object == f3.subject, and li1, li2, li3 all
    # distinct. Generated AFTER negatives with an independent, deterministically
    # seeded RNG so that adding this feature does not perturb the existing
    # recall / composition / negative question sequences, preserving all
    # committed test assertions for prior cycles.
    rng_3hop = random.Random(seed ^ 0xC0FFEE)
    chains_3hop: list[tuple[Fact, int, Fact, int, Fact, int]] = []
    for (f1, li1, f2, li2) in chains:
        for (f3, li3) in by_subject.get(f2.object, []):
            if li3 == li1 or li3 == li2:
                continue
            chains_3hop.append((f1, li1, f2, li2, f3, li3))
    rng_3hop.shuffle(chains_3hop)
    for (f1, li1, f2, li2, f3, li3) in chains_3hop[:n_compositions_3hop]:
        questions.append(Question(
            qid=qid, kind="composition_3hop",
            subject=f1.subject, relation=f1.relation,
            relation2=f2.relation, relation3=f3.relation,
            source_lessons=tuple(sorted({li1, li2, li3})),
            gold_object=f3.object,
        ))
        qid += 1

    return Dataset(
        vocab=vocab,
        lessons=lessons,
        questions=questions,
        entity_ids=entity_ids,
        relation_ids=relation_ids,
    )


def load_dataset_json(path: Path) -> Dataset:
    """Deserialize a `dataset.json` file produced by `train._serialize_dataset`.

    Used by `runs.replay` so that replay verifies against the exact
    dataset persisted at run time, independent of whether `data.py`
    has drifted since the run was recorded.
    """
    raw = json.loads(Path(path).read_text())
    # JSON stringifies integer dict keys; convert vocab keys back to ints.
    vocab = {int(k): v for k, v in raw["vocab"].items()}
    entity_ids = [int(x) for x in raw["entity_ids"]]
    relation_ids = [int(x) for x in raw["relation_ids"]]
    lessons = [
        Lesson(
            idx=int(L["idx"]),
            facts=[Fact(subject=int(s), relation=int(r), object=int(o))
                   for s, r, o in L["facts"]],
        )
        for L in raw["lessons"]
    ]
    questions = [
        Question(
            qid=int(q["qid"]),
            kind=q["kind"],
            subject=int(q["subject"]),
            relation=int(q["relation"]),
            relation2=int(q["relation2"]) if q.get("relation2") is not None else None,
            relation3=int(q["relation3"]) if q.get("relation3") is not None else None,
            candidate=int(q["candidate"]) if q.get("candidate") is not None else None,
            source_lessons=tuple(int(x) for x in q["source_lessons"]),
            gold_object=int(q["gold_object"]) if q.get("gold_object") is not None else None,
            gold_is_true=q.get("gold_is_true"),
        )
        for q in raw["questions"]
    ]
    return Dataset(
        vocab=vocab,
        lessons=lessons,
        questions=questions,
        entity_ids=entity_ids,
        relation_ids=relation_ids,
    )