tasks.py

"""Symbolic task definitions.

Everything a learning system sees — lessons, questions, answers — is expressed
as integer token IDs over a closed-world vocabulary. This keeps the evaluation
substrate language-agnostic and stops architectures from accidentally cheating
with surface-form tricks.
"""

from __future__ import annotations
from dataclasses import dataclass


@dataclass(frozen=True)
class Fact:
    subject: int
    relation: int
    object: int


@dataclass
class Lesson:
    idx: int
    facts: list[Fact]


@dataclass(frozen=True)
class Question:
    qid: int
    kind: str                            # "recall" | "composition" | "composition_3hop" | "negative"
    subject: int
    relation: int
    source_lessons: tuple[int, ...]      # lessons that taught the facts this Q depends on
    relation2: int | None = None         # composition / composition_3hop: second hop relation
    relation3: int | None = None         # composition_3hop: third hop relation
    candidate: int | None = None         # negative: proposed object to verify
    gold_object: int | None = None       # recall / composition / composition_3hop: expected object id
    gold_is_true: bool | None = None     # negative: is (s, r, candidate) a real taught fact?


@dataclass
class Answer:
    predicted_object: int | None = None
    predicted_is_true: bool | None = None


@dataclass
class Dataset:
    vocab: dict[int, str]                # token id -> invented word
    lessons: list[Lesson]
    questions: list[Question]
    entity_ids: list[int]
    relation_ids: list[int]