initium-deeplearning 0.1.0__py3-none-any.whl

initium/inference/tasks/__init__.py

@@ -0,0 +1 @@
+"""Inference tasks (one per dataset type), mirroring data/."""

initium/inference/tasks/audio_task.py

@@ -0,0 +1,114 @@
+"""
+file: inference/tasks/audio_task.py
+
+Audio-modality inference task. Delegates episode-building/decoding to
+data.audio.audio_dataset.AudioChainDataset (same pattern graph_traversal_task.py
+uses relative to graph_traversal.py - training and inference can never
+silently drift apart).
+"""
+
+from __future__ import annotations
+
+import torch
+
+from initium.data.audio.audio_dataset import AudioChainDataset
+from initium.inference.metrics import EpisodeScore
+from initium.inference.tasks.base_task import BaseInferenceTask, Episode
+
+BUILTIN_LINKS = (None, "", "audio-chain")
+
+
+class AudioChainTask(BaseInferenceTask):
+    name = "audio-chain"
+    dataset_type = "audio"
+
+    def __init__(self, dataset_link=None, path_length_range=None, **kwargs):
+        # Inference always tests a trained model, so --dataset-link here plays the
+        # role of test_dataset_link on the training-side dataset class: a real
+        # audio file used as the held-out fact source, instead of the synthetic
+        # seeded table. Omit (or pass one of BUILTIN_LINKS) for that default.
+        link = None if dataset_link in BUILTIN_LINKS else dataset_link
+        self._ds = AudioChainDataset(test_dataset_link=link)
+        self.input_dim, self.output_dim = self._ds.input_dim, self._ds.output_dim
+        self.query_range = (
+            tuple(path_length_range) if path_length_range else self._ds.ood_query_range
+        )
+        self.facts = self._ds.build_ood_facts()
+
+    def describe(self) -> str:
+        return f"{self.name} | {len(self.facts)} facts | queries {self.query_range}"
+
+    def build_episodes(self, n: int, rng, perturbation=None) -> list[Episode]:
+        severity = (
+            (perturbation or {}).get("severity", 0.0) if isinstance(perturbation, dict) else 0.0
+        )
+        keys, vals_map = [k for k, _ in self.facts], dict(self.facts)
+        NF = self._ds.codec.num_digits
+        episodes = []
+        for _ in range(n):
+            nq = rng.randint(*self.query_range)
+            query_keys = [rng.choice(keys) for _ in range(nq)]
+            facts = list(self.facts)
+            rng.shuffle(facts)
+            inputs, targets, mask = [], [], []
+            for i, (k, v) in enumerate(facts):
+                inputs.append(
+                    self._ds._encode_fact(
+                        k, v, 1.0 if i == 0 else 0.0, perturb=(severity or None), rng=rng
+                    )
+                )
+                targets.append([0] * (2 * NF))
+                mask.append(0)
+            for i, k in enumerate(query_keys):
+                inputs.append(self._ds._encode_query(k, 1.0 if i == 0 else 0.0))
+                targets.append([0] * (2 * NF))
+                mask.append(0)
+            cumsum = 0
+            for i, k in enumerate(query_keys):
+                cumsum = (cumsum + vals_map[k]) % self._ds.label_range
+                inputs.append(self._ds._encode_answer(1.0 if i == 0 else 0.0))
+                targets.append(
+                    self._ds.codec.label_to_digit_targets(vals_map[k])
+                    + self._ds.codec.label_to_digit_targets(cumsum)
+                )
+                mask.append(1)
+            episodes.append(
+                Episode(
+                    torch.stack(inputs),
+                    torch.tensor(targets, dtype=torch.long),
+                    torch.tensor(mask, dtype=torch.float32),
+                    meta={"num_queries": nq},
+                )
+            )
+        return episodes
+
+    def score_episode(self, output, episode: Episode, verbose: bool = False) -> EpisodeScore:
+        codec, D = self._ds.codec, self._ds.codec.label_dim
+        answer_idx = (episode.mask == 1).nonzero(as_tuple=True)[0]
+        f = {"value": [0, 0], "cumsum": [0, 0]}
+        n_correct = 0
+        for idx in answer_idx:
+            v_pred, c_pred = (
+                codec.decode_field(output[idx][0:D]),
+                codec.decode_field(output[idx][D : 2 * D]),
+            )
+            td = episode.target[idx].tolist()
+            nd = codec.num_digits
+            v_tgt = int("".join(map(str, td[0:nd])))
+            c_tgt = int("".join(map(str, td[nd : 2 * nd])))
+            v_ok, c_ok = v_pred == v_tgt, c_pred == c_tgt
+            f["value"][0] += int(v_ok)
+            f["value"][1] += 1
+            f["cumsum"][0] += int(c_ok)
+            f["cumsum"][1] += 1
+            n_correct += int(v_ok) + int(c_ok)
+            if verbose:
+                print(f"  value {v_pred}/{v_tgt} ok={v_ok} | cumsum {c_pred}/{c_tgt} ok={c_ok}")
+        n = 2 * len(answer_idx)
+        return EpisodeScore(
+            n_items=n,
+            n_correct=n_correct,
+            perfect=(n_correct == n),
+            group=len(answer_idx),
+            fields={k: (v[0], v[1]) for k, v in f.items()},
+        )

initium/inference/tasks/base_task.py

@@ -0,0 +1,60 @@
+"""
+file: inference/tasks/base_task.py
+
+Interface the inference engine talks to (counterpart of data/base_dataset.py).
+Any new type (text, audio, video, ...) implements BaseInferenceTask and
+registers itself in task_registry.py; the engine never imports anything
+task-specific.
+
+Contract:
+    dataset_type : canonical type string matched against the checkpoint's
+                   supported_dataset_types (e.g. "graph")
+    input_dim    : model input width the task produces (DNC input_size)
+    output_dim   : output_proj width the task's scoring expects
+    build_episodes(n, rng) -> list[Episode]
+    score_episode(output, episode, verbose) -> EpisodeScore
+        `output` is the post-output_proj (T, output_dim) CPU tensor.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import torch
+
+from initium.inference.metrics import EpisodeScore
+
+
+@dataclass
+class Episode:
+    input_seq: torch.Tensor  # (T, input_dim)
+    target: torch.Tensor  # task-defined, aligned with input_seq's time axis
+    mask: torch.Tensor  # (T,) 1 where a step is scored
+    meta: dict = field(default_factory=dict)
+    # Step indices where the model state may be snapshotted/reused (sorted). Everything
+    # before a boundary is a static prefix identical across episodes; no scored step
+    # (mask==1) may lie before a boundary. Empty -> nothing cacheable by the prefix cache.
+    cache_boundaries: list[int] = field(default_factory=list)
+
+
+class BaseInferenceTask:
+    name: str = "base"
+    dataset_type: str = "base"
+    input_dim: int | None = None
+    output_dim: int | None = None
+
+    def build_episodes(self, n: int, rng, perturbation=None) -> list[Episode]:
+        """perturbation: optional dataset-defined robustness probe (see
+        data/base_dataset.py's evaluate_robustness for the training-side
+        analogue), e.g. {'kind': 'noise', 'severity': 0.1}. None (default)
+        builds ordinary, unperturbed episodes -- purely additive, existing
+        tasks ignore it unless they opt in."""
+        raise NotImplementedError
+
+    def score_episode(
+        self, output: torch.Tensor, episode: Episode, verbose: bool = False
+    ) -> EpisodeScore:
+        raise NotImplementedError
+
+    def describe(self) -> str:
+        return self.name

initium/inference/tasks/graph_traversal_task.py

@@ -0,0 +1,235 @@
+"""
+file: inference/tasks/graph_traversal_task.py
+
+Graph-traversal inference task. Default test graph is the built-in London
+Underground (identical graph/label mapping to the training-time OOD eval).
+--dataset-link may instead point to your own edge file:
+
+  .csv / .tsv / .txt : one edge per row:  src,dst,line      ('#' comment lines and a
+                       header row starting with src/source/from are skipped)
+  .json              : [["src","dst","line"], ...]  or  {"edges": [...]}  (rows may also
+                       be {"src":..,"dst":..,"line":..})
+
+Names (stations/lines) are mapped to distinct 0-999 labels with a fixed
+random.Random(1234) exactly like the London eval. If EVERY field in the file
+is an integer in 0..999, the integers are used directly as labels instead.
+Episodes are built by the dataset module's own
+build_traversal_episode_from_graph(), so encoding is byte-identical to training.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from initium.data.common.graph_io import build_graph_from_raw_edges, load_raw_edges
+from initium.data.graph_traversal.graph_traversal import (
+    INPUT_DIM,
+    OOD_PATH_LENGTH_RANGE,
+    TRIPLE_DIM,
+    build_london_underground_eval,
+    build_traversal_episode_from_graph,
+    decode_prediction,
+    encode_triple,
+    triple_to_digit_targets,
+)
+from initium.inference.metrics import EpisodeScore
+from initium.inference.tasks.base_task import BaseInferenceTask, Episode
+
+BUILTIN_LINKS = (None, "", "graph-traversal", "london", "london-underground")
+
+
+class GraphTraversalTask(BaseInferenceTask):
+    name = "graph-traversal"
+    dataset_type = "graph"
+    input_dim = INPUT_DIM
+    output_dim = TRIPLE_DIM
+
+    def __init__(
+        self, dataset_link=None, path_length_range=None, shared_context=False, num_contexts=1
+    ):
+        if path_length_range is not None:
+            lo, hi = path_length_range
+            if lo < 1 or hi < lo:
+                raise ValueError(f"invalid path length range {path_length_range}")
+            self.path_length_range = (int(lo), int(hi))
+        else:
+            self.path_length_range = tuple(OOD_PATH_LENGTH_RANGE)
+
+        if int(num_contexts) < 1:
+            raise ValueError("num_contexts must be >= 1")
+        if int(num_contexts) > 1 and not shared_context:
+            raise ValueError("num_contexts > 1 requires shared_context=True")
+        self.shared_context, self.num_contexts = bool(shared_context), int(num_contexts)
+
+        if dataset_link in BUILTIN_LINKS:
+            self.edges, self.node_labels, self.adjacency = build_london_underground_eval()
+            self.source_desc = "London Underground (built-in)"
+        else:
+            # dataset_link may be a single edge file, a directory, a glob
+            # pattern, or a '+'-joined list of these -- load_raw_edges
+            # resolves it (data/common/real_data.py's resolve_link_paths)
+            # and raises its own clear FileNotFoundError if nothing
+            # matches, so no separate isfile guard is needed here.
+            self.edges, self.node_labels, self.adjacency = build_graph_from_raw_edges(
+                load_raw_edges(dataset_link)
+            )
+            self.source_desc = f"source: {dataset_link}"
+
+        if not any(self.adjacency[i] for i in self.adjacency):
+            raise ValueError("test graph has no outgoing edges; cannot build traversal episodes")
+
+    def describe(self) -> str:
+        shared = f" | shared context x{self.num_contexts}" if self.shared_context else ""
+        return (
+            f"{self.name} | {self.source_desc} | {len(self.node_labels)} nodes, "
+            f"{len(self.edges)} edges | walk length {self.path_length_range}{shared}"
+        )
+
+    def build_episodes(self, n: int, rng, perturbation=None) -> list[Episode]:
+        # perturbation: not wired for graph yet (see base_task.py) -- accepted
+        # and ignored so the shared BaseInferenceTask interface stays uniform.
+        episodes: list[Episode] = []
+        attempts = 0
+        if self.shared_context:
+            return self._build_shared_context_episodes(n, rng)
+        while len(episodes) < n:
+            attempts += 1
+            if attempts > 100 * n + 1000:
+                raise RuntimeError("could not build enough episodes from this graph")
+            ep = build_traversal_episode_from_graph(
+                self.edges,
+                self.node_labels,
+                self.adjacency,
+                len(self.node_labels),
+                self.path_length_range,
+                rng=rng,
+            )
+            if ep is None:
+                continue
+            input_seq, target_digits, answer_mask = ep
+            episodes.append(
+                Episode(
+                    input_seq,
+                    target_digits,
+                    answer_mask,
+                    meta={"walk_length": int(answer_mask.sum().item())},
+                )
+            )
+        return episodes
+
+        # ---- shared-context episodes (static prefix -> cacheable) ----------------
+
+    def _make_context(self, rng) -> torch.Tensor:
+        """One fixed, shuffled edge listing: the static prefix shared by many queries."""
+        shuffled = list(self.edges)
+        rng.shuffle(shuffled)
+        return torch.stack(
+            [
+                encode_triple(s, e, d, 1.0 if i == 0 else 0.0, 0.0)
+                for i, (s, e, d) in enumerate(shuffled)
+            ]
+        )
+
+    def _build_query(self, rng):
+        """Walk + answer steps only (mirrors the second half of
+        build_traversal_episode_from_graph). -> (inputs, target_digits, mask) or None."""
+        path_length = rng.randint(*self.path_length_range)
+        cur = rng.randrange(len(self.node_labels))
+        walk = []
+        for _ in range(path_length):
+            if not self.adjacency[cur]:
+                break
+            dst_idx, edge_label = rng.choice(self.adjacency[cur])
+            walk.append((cur, edge_label, dst_idx))
+            cur = dst_idx
+        if not walk:
+            return None
+        steps, targets, mask = [], [], []
+        for i, (src_idx, edge_label, _dst) in enumerate(walk):
+            src = self.node_labels[src_idx] if i == 0 else None
+            steps.append(encode_triple(src, edge_label, None, 1.0 if i == 0 else 0.0, 0.0))
+            targets.append([0] * 9)
+            mask.append(0)
+        for i, (src_idx, edge_label, dst_idx) in enumerate(walk):
+            steps.append(encode_triple(None, None, None, 1.0 if i == 0 else 0.0, 1.0))
+            targets.append(
+                triple_to_digit_targets(
+                    self.node_labels[src_idx], edge_label, self.node_labels[dst_idx]
+                )
+            )
+            mask.append(1)
+        return (
+            torch.stack(steps),
+            torch.tensor(targets, dtype=torch.long),
+            torch.tensor(mask, dtype=torch.float32),
+        )
+
+    def _build_shared_context_episodes(self, n: int, rng) -> list[Episode]:
+        contexts = [self._make_context(rng) for _ in range(self.num_contexts)]
+        episodes: list[Episode] = []
+        attempts = 0
+        while len(episodes) < n:
+            attempts += 1
+            if attempts > 100 * n + 1000:
+                raise RuntimeError("could not build enough episodes from this graph")
+            q = self._build_query(rng)
+            if q is None:
+                continue
+            k = len(episodes) % self.num_contexts
+            ctx_seq, (q_in, q_tgt, q_mask) = contexts[k], q
+            P = ctx_seq.shape[0]
+            episodes.append(
+                Episode(
+                    torch.cat([ctx_seq, q_in]),
+                    torch.cat([torch.zeros(P, 9, dtype=torch.long), q_tgt]),
+                    torch.cat([torch.zeros(P), q_mask]),
+                    meta={"walk_length": int(q_mask.sum().item()), "context_id": k},
+                    cache_boundaries=[P],
+                )
+            )
+        return episodes
+
+    def score_episode(self, output, episode: Episode, verbose: bool = False) -> EpisodeScore:
+        answer_idx = (episode.mask == 1).nonzero(as_tuple=True)[0]
+        f = {
+            "src": [0, 0],
+            "edge": [0, 0],
+            "dst": [0, 0],
+            "dst|src+edge": [0, 0],
+            "src|prev_dst_ok": [0, 0],
+        }
+        prev_dst_ok = None
+        n_correct = 0
+        for hop, idx in enumerate(answer_idx, start=1):
+            pred = decode_prediction(output[idx])
+            td = episode.target[idx].tolist()
+            tgt = (
+                int("".join(map(str, td[0:3]))),
+                int("".join(map(str, td[3:6]))),
+                int("".join(map(str, td[6:9]))),
+            )
+            ok = pred == tgt
+            n_correct += int(ok)
+            for name, p, t in zip(("src", "edge", "dst"), pred, tgt):
+                f[name][0] += int(p == t)
+                f[name][1] += 1
+            if pred[0] == tgt[0] and pred[1] == tgt[1]:  # lookup check
+                f["dst|src+edge"][1] += 1
+                f["dst|src+edge"][0] += int(pred[2] == tgt[2])
+            if prev_dst_ok:  # chain check
+                f["src|prev_dst_ok"][1] += 1
+                f["src|prev_dst_ok"][0] += int(pred[0] == tgt[0])
+            prev_dst_ok = pred[2] == tgt[2]
+            hp = f.setdefault(f"hop{hop}", [0, 0])  # triple acc per hop position
+            hp[0] += int(ok)
+            hp[1] += 1
+            if verbose:
+                print(f"  Pred: {pred} | Target: {tgt} | Correct: {ok}")
+        n = len(answer_idx)
+        return EpisodeScore(
+            n_items=n,
+            n_correct=n_correct,
+            perfect=(n_correct == n),
+            group=n,
+            fields={k: (v[0], v[1]) for k, v in f.items()},
+        )

initium/inference/tasks/multimodal_task.py

@@ -0,0 +1,74 @@
+"""
+file: inference/tasks/multimodal_task.py
+
+Multimodal inference task: delegates to data.multimodal.multimodal_dataset.
+MultimodalDataset the same way text/audio/video tasks delegate to their
+training-side dataset classes.
+"""
+
+from __future__ import annotations
+
+from initium.data.multimodal.multimodal_dataset import MultimodalDataset
+from initium.inference.metrics import EpisodeScore
+from initium.inference.tasks.base_task import BaseInferenceTask, Episode
+
+
+class MultimodalTask(BaseInferenceTask):
+    dataset_type = "multimodal"
+
+    def __init__(self, modalities, path_length_range=None, **kwargs):
+        # Inference always tests a trained model, so each per-modality real
+        # source given via --dataset-link (e.g. "video:clip.mp4+audio:clip.mp4")
+        # is wired in as that modality's test_dataset_link, matching every
+        # other *_task.py's own convention (see e.g. text_task.py).
+        self._ds = MultimodalDataset(modalities, link_role="test_dataset_link")
+        self.name = self._ds.name
+        self.input_dim, self.output_dim = self._ds.input_dim, self._ds.output_dim
+        self.query_range = (
+            tuple(path_length_range) if path_length_range else self._ds.primary.ood_query_range
+        )
+
+    def describe(self) -> str:
+        real = [
+            m for m, s in zip(self._ds.modalities, self._ds.subs) if s._test_fact_pool is not None
+        ]
+        src = f" | real data: {', '.join(real)}" if real else ""
+        return f"{self.name} | queries {self.query_range}{src}"
+
+    def build_episodes(self, n: int, rng, perturbation=None) -> list[Episode]:
+        episodes = []
+        for _ in range(n):
+            inp, tgt, mask, nq = self._ds._build_fused_episode(
+                (15, 20), self.query_range, rng=rng, use_test_pool=True
+            )
+            episodes.append(Episode(inp, tgt, mask, meta={"num_queries": nq}))
+        return episodes
+
+    def score_episode(self, output, episode: Episode, verbose: bool = False) -> EpisodeScore:
+        codec, D = self._ds.primary.codec, self._ds.primary.codec.label_dim
+        answer_idx = (episode.mask == 1).nonzero(as_tuple=True)[0]
+        f = {"value": [0, 0], "cumsum": [0, 0]}
+        n_correct = 0
+        for idx in answer_idx:
+            v_pred, c_pred = (
+                codec.decode_field(output[idx][0:D]),
+                codec.decode_field(output[idx][D : 2 * D]),
+            )
+            td = episode.target[idx].tolist()
+            nd = codec.num_digits
+            v_tgt = int("".join(map(str, td[0:nd])))
+            c_tgt = int("".join(map(str, td[nd : 2 * nd])))
+            v_ok, c_ok = v_pred == v_tgt, c_pred == c_tgt
+            f["value"][0] += int(v_ok)
+            f["value"][1] += 1
+            f["cumsum"][0] += int(c_ok)
+            f["cumsum"][1] += 1
+            n_correct += int(v_ok) + int(c_ok)
+        n = 2 * len(answer_idx)
+        return EpisodeScore(
+            n_items=n,
+            n_correct=n_correct,
+            perfect=(n_correct == n),
+            group=len(answer_idx),
+            fields={k: (v[0], v[1]) for k, v in f.items()},
+        )

initium/inference/tasks/task_registry.py

@@ -0,0 +1,68 @@
+"""
+file: inference/tasks/task_registry.py
+
+Single entry point used by run_inference.py (counterpart of
+data/dataset_registry.py):
+    task = get_task(dataset_type, dataset_link)
+
+v2: text/audio/video/multimodal, mirroring data/dataset_registry.py's v2.
+multimodal is excluded from IMPLEMENTED_TYPES (legacy dim-inference list in
+capabilities.py) since its input/output dims depend on WHICH modalities are
+combined, not a fixed per-type signature -- it's only usable as an explicit
+--dataset-type, never inferred from a legacy checkpoint's tensor shapes.
+"""
+
+GRAPH_ALIASES = ("graph", "graph-traversal", "graph_traversal")
+TEXT_ALIASES = ("text",)
+AUDIO_ALIASES = ("audio",)
+VIDEO_ALIASES = ("video",)
+MULTIMODAL_ALIASES = ("multimodal", "multi-modal", "multi_modal")
+IMPLEMENTED_TYPES = ("graph", "text", "audio", "video")
+
+
+def canonical_type(dataset_type: str) -> str:
+    t = (dataset_type or "graph").lower()
+    if t in GRAPH_ALIASES:
+        return "graph"
+    if t in TEXT_ALIASES:
+        return "text"
+    if t in AUDIO_ALIASES:
+        return "audio"
+    if t in VIDEO_ALIASES:
+        return "video"
+    if t in MULTIMODAL_ALIASES:
+        return "multimodal"
+    return t
+
+
+def get_task_class(dataset_type: str):
+    t = canonical_type(dataset_type)
+    if t == "graph":
+        from initium.inference.tasks.graph_traversal_task import GraphTraversalTask
+
+        return GraphTraversalTask
+    if t == "text":
+        from initium.inference.tasks.text_task import TextChainTask
+
+        return TextChainTask
+    if t == "audio":
+        from initium.inference.tasks.audio_task import AudioChainTask
+
+        return AudioChainTask
+    if t == "video":
+        from initium.inference.tasks.video_task import VideoChainTask
+
+        return VideoChainTask
+    if t == "multimodal":
+        from initium.inference.tasks.multimodal_task import MultimodalTask
+
+        return MultimodalTask
+    raise ValueError(f"unknown dataset type {dataset_type!r}")
+
+
+def get_task(dataset_type: str = "graph", dataset_link: str | None = None, **kwargs):
+    t = canonical_type(dataset_type)
+    cls = get_task_class(t)
+    if t == "multimodal":
+        return cls((dataset_link or "text+audio").split("+"), **kwargs)
+    return cls(dataset_link, **kwargs)

initium/inference/tasks/text_task.py

@@ -0,0 +1,114 @@
+"""
+file: inference/tasks/text_task.py
+
+Text-modality inference task. Delegates episode-building/decoding to
+data.text.text_dataset.TextChainDataset (same pattern graph_traversal_task.py
+uses relative to graph_traversal.py - training and inference can never
+silently drift apart).
+"""
+
+from __future__ import annotations
+
+import torch
+
+from initium.data.text.text_dataset import TextChainDataset
+from initium.inference.metrics import EpisodeScore
+from initium.inference.tasks.base_task import BaseInferenceTask, Episode
+
+BUILTIN_LINKS = (None, "", "text-chain")
+
+
+class TextChainTask(BaseInferenceTask):
+    name = "text-chain"
+    dataset_type = "text"
+
+    def __init__(self, dataset_link=None, path_length_range=None, **kwargs):
+        # Inference always tests a trained model, so --dataset-link here plays the
+        # role of test_dataset_link on the training-side dataset class: a real
+        # text file used as the held-out fact source, instead of the synthetic
+        # seeded table. Omit (or pass one of BUILTIN_LINKS) for that default.
+        link = None if dataset_link in BUILTIN_LINKS else dataset_link
+        self._ds = TextChainDataset(test_dataset_link=link)
+        self.input_dim, self.output_dim = self._ds.input_dim, self._ds.output_dim
+        self.query_range = (
+            tuple(path_length_range) if path_length_range else self._ds.ood_query_range
+        )
+        self.facts = self._ds.build_ood_facts()
+
+    def describe(self) -> str:
+        return f"{self.name} | {len(self.facts)} facts | queries {self.query_range}"
+
+    def build_episodes(self, n: int, rng, perturbation=None) -> list[Episode]:
+        severity = (
+            (perturbation or {}).get("severity", 0.0) if isinstance(perturbation, dict) else 0.0
+        )
+        keys, vals_map = [k for k, _ in self.facts], dict(self.facts)
+        NF = self._ds.codec.num_digits
+        episodes = []
+        for _ in range(n):
+            nq = rng.randint(*self.query_range)
+            query_keys = [rng.choice(keys) for _ in range(nq)]
+            facts = list(self.facts)
+            rng.shuffle(facts)
+            inputs, targets, mask = [], [], []
+            for i, (k, v) in enumerate(facts):
+                inputs.append(
+                    self._ds._encode_fact(
+                        k, v, 1.0 if i == 0 else 0.0, perturb=(severity or None), rng=rng
+                    )
+                )
+                targets.append([0] * (2 * NF))
+                mask.append(0)
+            for i, k in enumerate(query_keys):
+                inputs.append(self._ds._encode_query(k, 1.0 if i == 0 else 0.0))
+                targets.append([0] * (2 * NF))
+                mask.append(0)
+            cumsum = 0
+            for i, k in enumerate(query_keys):
+                cumsum = (cumsum + vals_map[k]) % self._ds.label_range
+                inputs.append(self._ds._encode_answer(1.0 if i == 0 else 0.0))
+                targets.append(
+                    self._ds.codec.label_to_digit_targets(vals_map[k])
+                    + self._ds.codec.label_to_digit_targets(cumsum)
+                )
+                mask.append(1)
+            episodes.append(
+                Episode(
+                    torch.stack(inputs),
+                    torch.tensor(targets, dtype=torch.long),
+                    torch.tensor(mask, dtype=torch.float32),
+                    meta={"num_queries": nq},
+                )
+            )
+        return episodes
+
+    def score_episode(self, output, episode: Episode, verbose: bool = False) -> EpisodeScore:
+        codec, D = self._ds.codec, self._ds.codec.label_dim
+        answer_idx = (episode.mask == 1).nonzero(as_tuple=True)[0]
+        f = {"value": [0, 0], "cumsum": [0, 0]}
+        n_correct = 0
+        for idx in answer_idx:
+            v_pred, c_pred = (
+                codec.decode_field(output[idx][0:D]),
+                codec.decode_field(output[idx][D : 2 * D]),
+            )
+            td = episode.target[idx].tolist()
+            nd = codec.num_digits
+            v_tgt = int("".join(map(str, td[0:nd])))
+            c_tgt = int("".join(map(str, td[nd : 2 * nd])))
+            v_ok, c_ok = v_pred == v_tgt, c_pred == c_tgt
+            f["value"][0] += int(v_ok)
+            f["value"][1] += 1
+            f["cumsum"][0] += int(c_ok)
+            f["cumsum"][1] += 1
+            n_correct += int(v_ok) + int(c_ok)
+            if verbose:
+                print(f"  value {v_pred}/{v_tgt} ok={v_ok} | cumsum {c_pred}/{c_tgt} ok={c_ok}")
+        n = 2 * len(answer_idx)
+        return EpisodeScore(
+            n_items=n,
+            n_correct=n_correct,
+            perfect=(n_correct == n),
+            group=len(answer_idx),
+            fields={k: (v[0], v[1]) for k, v in f.items()},
+        )