tinymind-native-8b-remote-handoff/bundle/evaluation/pure_lattice_cnn.py

"""Evidence report for TinyMind PureLattice CNN core."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import torch

from model.architecture import OmegaModel
from model.config import OmegaConfig
from model.pure_lattice_cnn import PureLatticeCNNConfig, PureLatticeCNNCore, count_trainable_parameters


def build_pure_lattice_cnn_report(
    out_dir: str | Path,
    *,
    dim: int = 64,
    seq_len: int = 33,
    batch_size: int = 2,
    seed: int = 20260526,
) -> dict[str, Any]:
    torch.manual_seed(seed)
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    report_path = out / "pure_lattice_cnn_report.json"

    cfg = PureLatticeCNNConfig(dim=dim)
    core = PureLatticeCNNCore(cfg)
    x = torch.randn(batch_size, seq_len, dim)
    y, token_state = core(x)
    forward_finite = bool(torch.isfinite(y).all().item())

    loss = y.float().pow(2).mean()
    loss.backward()
    backward_finite = all(
        param.grad is None or bool(torch.isfinite(param.grad).all().item())
        for param in core.parameters()
    )

    grid = torch.randn(batch_size, dim, 4, max(2, seq_len // 4))
    grid_y, grid_state = core.forward_grid(grid)
    grid_finite = bool(torch.isfinite(grid_y).all().item())

    omega_cfg = OmegaConfig(
        vocab_size=128,
        dim=dim,
        n_layers=1,
        n_heads=max(1, dim // 16),
        head_dim=16 if dim >= 16 else dim,
        ffn_mult=2,
        dropout=0.0,
        cnn_core_enabled=True,
    )
    omega = OmegaModel(omega_cfg)
    input_ids = torch.randint(4, omega_cfg.vocab_size, (batch_size, min(seq_len, 17)))
    omega_out = omega(input_ids, labels=input_ids)
    omega_finite = bool(torch.isfinite(omega_out["logits"]).all().item() and torch.isfinite(omega_out["loss"]).item())

    report: dict[str, Any] = {
        "schema_version": "tinymind-pure-lattice-cnn-v1",
        "report_path": str(report_path),
        "config": {
            "dim": cfg.dim,
            "hidden_mult": cfg.hidden_mult,
            "kernel_sizes": list(cfg.kernel_sizes),
            "dilations": list(cfg.dilations),
            "dropout": cfg.dropout,
            "residual_scale": cfg.residual_scale,
            "receptive_field": core.receptive_field,
        },
        "parameter_count": count_trainable_parameters(core),
        "token_probe": token_state,
        "grid_probe": grid_state,
        "forward_finite": forward_finite,
        "backward_finite": backward_finite,
        "grid_forward_finite": grid_finite,
        "omega_integration_probe": {
            "enabled": omega.cnn_stem is not None,
            "forward_finite": omega_finite,
            "logits_shape": list(omega_out["logits"].shape),
        },
        "claim_gate": {
            "cnn_core_ready": forward_finite and backward_finite and grid_finite,
            "integrated_into_omega_model": bool(omega.cnn_stem is not None and omega_finite),
            "integrated_into_12b_runtime": False,
            "native_multimodal_claim_allowed": False,
            "world_best_cnn_claim_allowed": False,
            "requires_downstream_training_evidence": True,
        },
    }
    report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    return report