initium-deeplearning 0.1.0__py3-none-any.whl

initium/inference/cache/stores.py

@@ -0,0 +1,167 @@
+"""
+file: inference/cache/stores.py
+
+Storage back-ends shared by all cache types: a byte-capped RAM LRU, an optional
+on-disk store (survives across CLI invocations), and a tiered wrapper. Entries are
+plain dicts {"payload", "compute_ms", "nbytes", "kind"} made only of tensors /
+lists / tuples / dicts / numbers, so the disk tier can use torch.load(weights_only=True)
+(no arbitrary pickle execution).
+"""
+
+from __future__ import annotations
+
+import glob
+import os
+from collections import OrderedDict
+from dataclasses import asdict, dataclass
+
+import torch
+
+
+@dataclass
+class CacheStats:
+    lookups: int = 0
+    hits: int = 0
+    misses: int = 0
+    stores: int = 0
+    ram_hits: int = 0
+    disk_hits: int = 0
+    bytes_stored: int = 0
+    est_saved_ms: float = 0.0  # sum of recorded cold-compute time of every reused entry
+
+    def hit(self, tier: str, saved_ms: float) -> None:
+        self.hits += 1
+        self.est_saved_ms += float(saved_ms)
+        if tier == "disk":
+            self.disk_hits += 1
+        else:
+            self.ram_hits += 1
+
+    def to_dict(self) -> dict:
+        d = asdict(self)
+        d["hit_rate"] = 100.0 * self.hits / max(self.lookups, 1)
+        return d
+
+
+class MemoryStore:
+    def __init__(self, max_bytes: int):
+        self.max_bytes = max_bytes
+        self._d: OrderedDict[str, dict] = OrderedDict()
+        self.bytes = 0
+        self.evictions = 0
+
+    def get(self, key: str) -> dict | None:
+        e = self._d.get(key)
+        if e is not None:
+            self._d.move_to_end(key)
+        return e
+
+    def put(self, key: str, entry: dict) -> None:
+        nb = int(entry.get("nbytes", 0))
+        if nb > self.max_bytes:
+            return
+        old = self._d.pop(key, None)
+        if old is not None:
+            self.bytes -= int(old.get("nbytes", 0))
+        self._d[key] = entry
+        self.bytes += nb
+        while self.bytes > self.max_bytes and self._d:
+            _, ev = self._d.popitem(last=False)
+            self.bytes -= int(ev.get("nbytes", 0))
+            self.evictions += 1
+
+    def __len__(self) -> int:
+        return len(self._d)
+
+
+class DiskStore:
+    def __init__(self, root: str, max_bytes: int):
+        self.root, self.max_bytes = root, max_bytes
+        os.makedirs(root, exist_ok=True)
+
+    def _path(self, key: str) -> str:
+        return os.path.join(self.root, key + ".pt")
+
+    def get(self, key: str) -> dict | None:
+        p = self._path(key)
+        if not os.path.isfile(p):
+            return None
+        try:
+            entry = torch.load(p, map_location="cpu", weights_only=True)
+            os.utime(p, None)  # recency for eviction
+            return entry
+        except Exception:  # corrupt / incompatible -> treat as miss
+            try:
+                os.remove(p)
+            except OSError:
+                pass
+            return None
+
+    def put(self, key: str, entry: dict) -> None:
+        p = self._path(key)
+        tmp = f"{p}.tmp{os.getpid()}"
+        try:
+            torch.save(entry, tmp)
+            os.replace(tmp, p)  # atomic
+        except Exception as e:
+            print(f"[cache] disk write failed for {key}: {e}")
+            try:
+                os.remove(tmp)
+            except OSError:
+                pass
+            return
+        self._evict()
+
+    def _evict(self) -> None:
+        files = [
+            (os.path.getmtime(f), os.path.getsize(f), f)
+            for f in glob.glob(os.path.join(self.root, "*.pt"))
+        ]
+        total = sum(s for _, s, _ in files)
+        for _, size, f in sorted(files):
+            if total <= self.max_bytes:
+                break
+            try:
+                os.remove(f)
+                total -= size
+            except OSError:
+                pass
+
+    def clear(self) -> None:
+        for f in glob.glob(os.path.join(self.root, "*.pt")):
+            try:
+                os.remove(f)
+            except OSError:
+                pass
+
+
+class TieredStore:
+    """RAM first, disk second (disk hits are promoted to RAM). Shared by every cache
+    of a run, so --cache-ram-mb / --cache-disk-mb are TOTAL budgets."""
+
+    def __init__(self, mem: MemoryStore, disk: DiskStore | None = None):
+        self.mem, self.disk = mem, disk
+
+    def get(self, key: str) -> tuple[dict, str] | None:
+        e = self.mem.get(key)
+        if e is not None:
+            return e, "ram"
+        if self.disk is not None:
+            e = self.disk.get(key)
+            if e is not None:
+                self.mem.put(key, e)
+                return e, "disk"
+        return None
+
+    def put(self, key: str, entry: dict) -> None:
+        self.mem.put(key, entry)
+        if self.disk is not None:
+            self.disk.put(key, entry)
+
+    def info(self) -> dict:
+        return {
+            "ram_mb": self.mem.bytes / 1e6,
+            "ram_entries": len(self.mem),
+            "ram_evictions": self.mem.evictions,
+            "disk_dir": self.disk.root if self.disk else None,
+        }

initium/inference/capabilities.py

@@ -0,0 +1,86 @@
+"""
+file: inference/capabilities.py
+
+Decides, FROM THE CHECKPOINT ITSELF, which dataset types a model can run, and
+refuses to continue when the requested type/dimensions don't match.
+
+Sources of truth (in order):
+  1. model_config["supported_dataset_types"] (written by core_training v19+).
+     May contain "*" = any type whose dims match the weights.
+  2. Legacy checkpoints (no such key): infer from the tensor dims - every
+     implemented task type whose (input_dim, output_dim) equals the model's.
+     Dims that match nothing -> supported list is empty -> run refuses.
+Independently of (1)/(2), the task's dims must ALWAYS equal the model's dims.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from initium.inference.inference_config import WILDCARD_TYPE
+from initium.inference.tasks.task_registry import IMPLEMENTED_TYPES, canonical_type, get_task_class
+
+
+class IncompatibleModelError(Exception):
+    pass
+
+
+@dataclass
+class ModelCapabilities:
+    supported_types: list[str]
+    input_dim: int
+    output_dim: int
+    source: str  # "recorded" | "inferred-legacy" | "unknown"
+    controller_type: str
+
+    def supports(self, canon_type: str) -> bool:
+        return WILDCARD_TYPE in self.supported_types or canon_type in self.supported_types
+
+    def describe(self) -> str:
+        return (
+            f"supported types={self.supported_types} (source: {self.source}) | "
+            f"input_dim={self.input_dim} output_dim={self.output_dim} | "
+            f"controller={self.controller_type}"
+        )
+
+
+def detect_capabilities(ckpt: dict) -> ModelCapabilities:
+    cfg = ckpt["model_config"]
+    out_dim, proj_in = ckpt["output_proj_state_dict"]["weight"].shape
+    input_dim = int(cfg.get("input_dim", cfg.get("input_size", proj_in)))
+    output_dim = int(out_dim)
+
+    recorded = cfg.get("supported_dataset_types")
+    if recorded:
+        types = [WILDCARD_TYPE if t == WILDCARD_TYPE else canonical_type(t) for t in recorded]
+        source = "recorded"
+    else:
+        types = []
+        for t in IMPLEMENTED_TYPES:
+            cls = get_task_class(t)
+            if (cls.input_dim, cls.output_dim) == (input_dim, output_dim):
+                types.append(t)
+        source = "inferred-legacy" if types else "unknown"
+
+    return ModelCapabilities(
+        types, input_dim, output_dim, source, cfg.get("controller_type", "lstm")
+    )
+
+
+def require_type_supported(caps: ModelCapabilities, dataset_type: str, ckpt_path: str) -> str:
+    canon = canonical_type(dataset_type)
+    if not caps.supports(canon):
+        raise IncompatibleModelError(
+            f"checkpoint '{ckpt_path}' cannot run dataset type '{canon}'. "
+            f"It supports {caps.supported_types} ({caps.source}). Not continuing."
+        )
+    return canon
+
+
+def require_dims_match(caps: ModelCapabilities, task, ckpt_path: str) -> None:
+    if (task.input_dim, task.output_dim) != (caps.input_dim, caps.output_dim):
+        raise IncompatibleModelError(
+            f"task '{task.name}' produces input_dim={task.input_dim}/output_dim={task.output_dim} "
+            f"but checkpoint '{ckpt_path}' expects input_dim={caps.input_dim}/"
+            f"output_dim={caps.output_dim}. Not continuing."
+        )

initium/inference/checkpoint_io.py

@@ -0,0 +1,63 @@
+"""
+file: inference/checkpoint_io.py
+
+Load a checkpoint written by core_training.save_checkpoint() (periodic or
+final - same format) and describe it. Loads to CPU; weights_only=False for the
+same reason as core_training.load_checkpoint_for_resume (self-produced file
+that also stores python/numpy RNG state).
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+
+REQUIRED_KEYS = ("rnn_state_dict", "output_proj_state_dict", "model_config")
+
+
+def load_checkpoint(path: str, model_config_override: dict | None = None) -> dict:
+    if not os.path.isfile(path):
+        raise FileNotFoundError(f"checkpoint not found: {path}")
+    ckpt = torch.load(path, map_location="cpu", weights_only=False)
+    if not isinstance(ckpt, dict):
+        raise ValueError(f"{path}: not a core_training checkpoint (expected a dict)")
+    missing = [k for k in REQUIRED_KEYS if k not in ckpt]
+    if missing:
+        raise ValueError(f"{path}: not a usable checkpoint, missing keys {missing}")
+    cfg = dict(ckpt["model_config"])
+    if model_config_override:
+        cfg.update(model_config_override)
+        print(f"[checkpoint_io] model_config overridden: {model_config_override}")
+    ckpt["model_config"] = cfg
+    return ckpt
+
+
+def describe_checkpoint(ckpt: dict) -> str:
+    c = ckpt["model_config"]
+    lines = [
+        f"run_id={ckpt.get('run_id')} step={ckpt.get('step')} beta_target={ckpt.get('beta_target')}",
+        f"controller={c.get('controller_type', 'lstm')} hidden={c.get('hidden_size')} "
+        f"nr_cells={c.get('nr_cells')} cell_size={c.get('cell_size')} read_heads={c.get('read_heads')}",
+        f"link_matrix={c.get('link_matrix_mode', 'dense')} moe={c.get('moe_enabled', False)} "
+        f"split_graph={c.get('split_graph_enabled', False)}"
+        + (f"({c.get('split_graph_variant')})" if c.get("split_graph_enabled") else "")
+        + f" dynamic_n={c.get('dynamic_n_mode', False)}",
+    ]
+    combiner_variant = c.get("split_graph_combiner_variant")
+    if (
+        c.get("split_graph_enabled")
+        and c.get("split_graph_combiner_mode") == "controller"
+        and combiner_variant
+    ):
+        if "+" in combiner_variant:
+            lines.append(
+                f"combiner=controller stages={list(enumerate(combiner_variant.split('+')))} "
+                f"(use --ablate-combiner-stage <name|index>)"
+            )
+        else:
+            lines.append(
+                f"combiner=controller variant={combiner_variant!r} (single stage -- "
+                f"multi-source MoE routing applies only here, not to hybrid chains)"
+            )
+    return "\n".join(lines)

initium/inference/engine.py

@@ -0,0 +1,88 @@
+"""
+file: inference/engine.py
+
+Runs episodes through a loaded model. Knows nothing about any specific task;
+it only calls task.score_episode().
+
+Two regimes (the reset_experience switch):
+  reset_experience=True  : every episode starts from hidden=(None,None,None) with
+                           reset_experience=True -> independent episodes, exactly
+                           like the periodic OOD eval in training.
+  reset_experience=False : the (controller state, memory, last read) returned by
+                           episode i is fed into episode i+1 with
+                           reset_experience=False -> memory content and recurrent
+                           state persist across the whole run. NO weights change;
+                           any "learning" is in-context via memory + controller state.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import torch
+
+from initium.inference.metrics import EpisodeResult
+
+FRESH_HIDDEN = (None, None, None)
+
+
+class InferenceEngine:
+    def __init__(
+        self,
+        model,
+        output_proj,
+        task,
+        device,
+        ablate_memory: bool = False,
+        combiner_skip_stages: frozenset | None = None,
+    ):
+        self.model, self.output_proj, self.task = model, output_proj, task
+        self.device, self.ablate_memory = device, ablate_memory
+        self.combiner_skip_stages = combiner_skip_stages
+        self.model.eval()
+        self.output_proj.eval()
+
+    @torch.no_grad()
+    def _forward(self, episode, hidden, reset_experience: bool):
+        x = episode.input_seq.unsqueeze(0).to(self.device)  # (1, T, input_dim)
+        kw: dict[str, Any] = dict(
+            reset_experience=reset_experience, pass_through_memory=not self.ablate_memory
+        )
+        if self.combiner_skip_stages:
+            kw["combiner_skip_stages"] = self.combiner_skip_stages
+        output, new_hidden = self.model(x, hidden, **kw)
+        output = output.transpose(0, 1).contiguous().squeeze(0)  # (T, input_dim)
+        output = self.output_proj(output)  # (T, output_dim)
+        return output.float().cpu(), new_hidden
+
+    def run(
+        self,
+        episodes,
+        reset_experience: bool,
+        verbose_n: int = 0,
+        progress_every: int = 0,
+        label: str = "",
+    ) -> list[EpisodeResult]:
+        hidden = FRESH_HIDDEN
+        results: list[EpisodeResult] = []
+        run_items = run_correct = 0
+        for i, ep in enumerate(episodes):
+            if reset_experience:
+                hidden = FRESH_HIDDEN
+            t0 = time.perf_counter()
+            out, hidden = self._forward(ep, hidden, reset_experience)
+            if self.device.type == "cuda":
+                torch.cuda.synchronize(self.device)
+            elapsed_ms = (time.perf_counter() - t0) * 1000.0
+
+            score = self.task.score_episode(out, ep, verbose=(i < verbose_n))
+            results.append(EpisodeResult(i, score, elapsed_ms, reset_experience))
+            run_items += score.n_items
+            run_correct += score.n_correct
+            if progress_every and ((i + 1) % progress_every == 0 or i + 1 == len(episodes)):
+                print(
+                    f"[engine{label}] episode {i + 1}/{len(episodes)} | "
+                    f"running item acc {100.0 * run_correct / max(run_items, 1):.2f}%"
+                )
+        return results

initium/inference/inference_config.py

@@ -0,0 +1,22 @@
+"""
+file: inference/inference_config.py
+
+Inference-side constants. Pure constants, imported explicitly by the other
+inference modules (same convention as config/train_config.py).
+"""
+
+INFERENCE_LOG_DIR = "./inference_logs"
+
+# Default task: graph traversal on the London Underground graph.
+DEFAULT_DATASET_TYPE = "graph"
+DEFAULT_DATASET_LINK = None  # None -> built-in London Underground graph
+
+DEFAULT_NUM_EPISODES = 200  # fresh mode (and persistent mode when --loop is omitted)
+DEFAULT_WINDOW = 10  # episodes per window in the persistent-mode adaptation table
+DEFAULT_SEED = 0
+DEFAULT_VERBOSE_N = 3  # print predictions of the first N episodes
+PROGRESS_EVERY = 25  # running-accuracy console line every N episodes (0 = off)
+
+# A checkpoint whose supported_dataset_types contains this accepts ANY task
+# whose input/output dims match its weights.
+WILDCARD_TYPE = "*"

initium/inference/metrics.py

@@ -0,0 +1,103 @@
+"""
+file: inference/metrics.py
+
+Task-agnostic scoring containers and aggregation. A task's score_episode()
+returns an EpisodeScore; everything downstream (breakdowns, windows,
+adaptation trend, logging) only ever sees EpisodeScore, so a new task type
+needs no changes here.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+
+
+@dataclass
+class EpisodeScore:
+    n_items: int  # scored items in the episode (graph: answer triples)
+    n_correct: int
+    perfect: bool  # every item correct
+    group: Any = None  # bucket key for breakdowns (graph: walk length / hop count)
+    fields: dict[str, tuple[int, int]] = field(default_factory=dict)  # name -> (correct, total)
+
+    @property
+    def item_acc(self) -> float:
+        return 100.0 * self.n_correct / max(self.n_items, 1)
+
+
+@dataclass
+class EpisodeResult:
+    index: int
+    score: EpisodeScore
+    elapsed_ms: float
+    reset_experience: bool
+    cache_event: str = ""  # "", "miss", "hit:result", "hit:prefix@22", ...
+
+
+def aggregate(scores: list[EpisodeScore]) -> dict:
+    n_eps = len(scores)
+    n_items = sum(s.n_items for s in scores)
+    n_correct = sum(s.n_correct for s in scores)
+    n_perfect = sum(int(s.perfect) for s in scores)
+
+    by_group: dict[Any, list[int]] = {}
+    field_tot: dict[str, list[int]] = {}
+    for s in scores:
+        if s.group is not None:
+            g = by_group.setdefault(s.group, [0, 0, 0, 0])
+            g[0] += s.n_items
+            g[1] += s.n_correct
+            g[2] += 1
+            g[3] += int(s.perfect)
+        for name, (c, t) in s.fields.items():
+            f = field_tot.setdefault(name, [0, 0])
+            f[0] += c
+            f[1] += t
+
+    return {
+        "n_episodes": n_eps,
+        "item_acc": 100.0 * n_correct / max(n_items, 1),
+        "perfect_frac": 100.0 * n_perfect / max(n_eps, 1),
+        "by_group": {
+            g: {
+                "item_acc": 100.0 * v[1] / max(v[0], 1),
+                "perfect_frac": 100.0 * v[3] / max(v[2], 1),
+                "n_episodes": v[2],
+            }
+            for g, v in sorted(by_group.items())
+        },
+        "field_acc": {n: 100.0 * c / max(t, 1) for n, (c, t) in field_tot.items()},
+    }
+
+
+def windowed(scores: list[EpisodeScore], window: int) -> list[dict]:
+    """Aggregate consecutive chunks of `window` episodes (persistent-mode adaptation view)."""
+    window = max(1, window)
+    out = []
+    for w, start in enumerate(range(0, len(scores), window)):
+        chunk = scores[start : start + window]
+        agg = aggregate(chunk)
+        agg.update(window_index=w, start_episode=start, end_episode=start + len(chunk) - 1)
+        out.append(agg)
+    return out
+
+
+def adaptation_trend(scores: list[EpisodeScore], window: int) -> dict:
+    """First-window vs last-window accuracy and a least-squares slope of
+    per-episode accuracy -- a compact 'did it adapt mid-inference' read."""
+    if len(scores) < 2:
+        return {}
+    w = max(1, min(window, len(scores) // 2))
+    first, last = aggregate(scores[:w]), aggregate(scores[-w:])
+    accs = np.array([s.item_acc for s in scores], dtype=np.float64)
+    slope = float(np.polyfit(np.arange(len(accs)), accs, 1)[0])
+    return {
+        "window": w,
+        "first_window_item_acc": first["item_acc"],
+        "last_window_item_acc": last["item_acc"],
+        "delta_item_acc": last["item_acc"] - first["item_acc"],
+        "slope_acc_per_episode": slope,
+    }