tinymind-native-colab-handoff/bundle/model/native_axiom_regenesis.py

from __future__ import annotations

from dataclasses import asdict, dataclass
from datetime import datetime, timezone
import json
import math
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F

from .axiom_dim import AxiomDimBridge, AxiomDimConfig
from .layers import RMSNorm
from .self_assessment_core import SelfAssessmentCore, SelfAssessmentCoreConfig


@dataclass
class AxiomReGenesisConfig:
    architecture_name: str = "TinyMind-AxiomReGenesis"
    vocab_size: int = 512
    dim: int = 512
    n_layers: int = 12
    lanes: int = 16
    max_seq_len: int = 512
    local_window: int = 64
    memory_slots: int = 8
    memory_rank: int = 32
    regen_top_k: int = 4
    regen_rank: int = 4
    axiom_effective_dim: int = 20_480
    axiom_basis_rank: int = 64
    axiom_facets: int = 8
    self_assessment_steps: int = 2
    dropout: float = 0.0
    residual_alpha: float | None = None


class NativeReGenesisKVLayer(nn.Module):
    """Regenerate compact K/V from retrieved exact chunks without storing historical KV."""

    def __init__(self, cfg: AxiomReGenesisConfig):
        super().__init__()
        self.cfg = cfg
        self.token_embed = nn.Embedding(cfg.vocab_size, cfg.dim)
        self.query_norm = RMSNorm(cfg.dim)
        self.chunk_norm = RMSNorm(cfg.dim)
        self.k_proj = nn.Linear(cfg.dim, cfg.regen_rank * cfg.dim, bias=False)
        self.v_proj = nn.Linear(cfg.dim, cfg.regen_rank * cfg.dim, bias=False)
        self.fuse = nn.Linear(cfg.dim * 2, cfg.dim, bias=False)
        self.gate = nn.Linear(cfg.dim, cfg.dim)

    def _normalize_tokens(self, retrieved_tokens: torch.Tensor | None, batch: int, device: torch.device) -> torch.Tensor:
        if retrieved_tokens is None:
            return torch.zeros(batch, self.cfg.regen_top_k, 1, device=device, dtype=torch.long)
        tokens = retrieved_tokens.to(device=device, dtype=torch.long).clamp(0, self.cfg.vocab_size - 1)
        if tokens.dim() == 2:
            tokens = tokens.unsqueeze(0).expand(batch, -1, -1).contiguous()
        if tokens.shape[1] > self.cfg.regen_top_k:
            tokens = tokens[:, : self.cfg.regen_top_k]
        elif tokens.shape[1] < self.cfg.regen_top_k:
            pad = torch.zeros(
                batch,
                self.cfg.regen_top_k - tokens.shape[1],
                tokens.shape[2],
                device=device,
                dtype=torch.long,
            )
            tokens = torch.cat([tokens, pad], dim=1)
        return tokens

    def forward(self, hidden: torch.Tensor, retrieved_tokens: torch.Tensor | None = None) -> tuple[torch.Tensor, dict[str, Any]]:
        batch, seq_len, _ = hidden.shape
        tokens = self._normalize_tokens(retrieved_tokens, batch, hidden.device)
        chunks = self.chunk_norm(self.token_embed(tokens).mean(dim=2))
        regen_k = self.k_proj(chunks).view(batch, self.cfg.regen_top_k, self.cfg.regen_rank, self.cfg.dim)
        regen_v = self.v_proj(chunks).view(batch, self.cfg.regen_top_k, self.cfg.regen_rank, self.cfg.dim)

        query = self.query_norm(hidden)
        q = F.normalize(query.mean(dim=1), dim=-1, eps=1e-6)
        k = F.normalize(regen_k.mean(dim=2), dim=-1, eps=1e-6)
        weights = torch.softmax(torch.einsum("bd,bkd->bk", q, k), dim=-1)
        value = torch.einsum("bk,bkrd->bd", weights, regen_v).unsqueeze(1).expand(batch, seq_len, self.cfg.dim)
        delta = torch.tanh(self.fuse(torch.cat([query, value], dim=-1)))
        out = hidden + torch.sigmoid(self.gate(query)) * delta
        weights_cpu = weights.detach().float().cpu()
        return out, {
            "regen_k_shape": list(regen_k.shape),
            "regen_v_shape": list(regen_v.shape),
            "retrieval_weights": weights_cpu.tolist(),
            "retrieval_weight_max": float(weights_cpu.max().item()),
            "retrieval_weight_min": float(weights_cpu.min().item()),
            "kv_tokens_stored": 0,
        }


class AxiomReGenesisBlock(nn.Module):
    """TinyMind-native factorized recurrent block with local exact and regenerated KV lanes."""

    def __init__(self, cfg: AxiomReGenesisConfig, layer_index: int):
        super().__init__()
        self.cfg = cfg
        self.layer_index = layer_index
        self.norm = RMSNorm(cfg.dim)
        self.axiom = AxiomDimBridge(
            AxiomDimConfig(
                physical_dim=cfg.dim,
                effective_dim=cfg.axiom_effective_dim,
                basis_rank=cfg.axiom_basis_rank,
                facets=cfg.axiom_facets,
                residual_scale=0.12,
            )
        )
        self.local_q = nn.Linear(cfg.dim, cfg.dim, bias=False)
        self.local_k = nn.Linear(cfg.dim, cfg.dim, bias=False)
        self.local_v = nn.Linear(cfg.dim, cfg.dim, bias=False)
        self.local_o = nn.Linear(cfg.dim, cfg.dim, bias=False)
        self.write_gate = nn.Linear(cfg.dim, cfg.memory_slots * cfg.memory_rank)
        self.value = nn.Linear(cfg.dim, cfg.memory_slots * cfg.memory_rank)
        self.read_gate = nn.Linear(cfg.dim, cfg.memory_slots)
        self.memory_out = nn.Linear(cfg.memory_rank, cfg.dim, bias=False)
        self.router = nn.Linear(cfg.dim, cfg.lanes)
        self.lane_proj = nn.Linear(cfg.lanes, cfg.dim, bias=False)
        self.regen = NativeReGenesisKVLayer(cfg)
        self.ffn = nn.Sequential(
            nn.Linear(cfg.dim, cfg.dim * 2),
            nn.SiLU(),
            nn.Dropout(cfg.dropout),
            nn.Linear(cfg.dim * 2, cfg.dim),
        )
        self.out_norm = RMSNorm(cfg.dim)

    def _local_exact_window(self, u: torch.Tensor) -> tuple[torch.Tensor, dict[str, Any]]:
        batch, seq_len, dim = u.shape
        q = self.local_q(u)
        k = self.local_k(u)
        v = self.local_v(u)
        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(max(1, dim))
        positions = torch.arange(seq_len, device=u.device)
        causal = positions[None, :] <= positions[:, None]
        window = (positions[:, None] - positions[None, :]) < max(1, int(self.cfg.local_window))
        mask = causal & window
        scores = scores.masked_fill(~mask.unsqueeze(0), torch.finfo(scores.dtype).min)
        attn = torch.softmax(scores, dim=-1)
        local = self.local_o(torch.matmul(attn, v))
        return local, {
            "local_window": int(self.cfg.local_window),
            "exact_recent_tokens": int(min(seq_len, self.cfg.local_window)),
            "attention_shape": [batch, seq_len, seq_len],
        }

    def _memory(self, u: torch.Tensor, state: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
        batch, seq_len, _ = u.shape
        if state is None:
            state = torch.zeros(batch, self.cfg.memory_slots, self.cfg.memory_rank, device=u.device, dtype=u.dtype)
        write = torch.sigmoid(self.write_gate(u)).view(batch, seq_len, self.cfg.memory_slots, self.cfg.memory_rank)
        value = torch.tanh(self.value(u)).view(batch, seq_len, self.cfg.memory_slots, self.cfg.memory_rank)
        # Contractive bounded update: average sequence writes into fixed memory slots.
        new_state = 0.92 * state + (write * value).mean(dim=1)
        new_state = torch.tanh(new_state)
        read = torch.softmax(self.read_gate(u), dim=-1)
        memory = torch.einsum("bts,bsr->btr", read, new_state)
        memory_hidden = self.memory_out(memory)
        return memory_hidden, new_state, {
            "memory_state_shape": list(new_state.shape),
            "memory_state_norm": float(new_state.detach().float().norm(dim=-1).mean().cpu()),
        }

    def forward(
        self,
        hidden: torch.Tensor,
        state: torch.Tensor | None = None,
        retrieved_tokens: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
        u = self.norm(hidden)
        axiom = self.axiom(u)
        local, local_report = self._local_exact_window(u)
        memory, next_state, memory_report = self._memory(u, state)
        routed = self.lane_proj(torch.softmax(self.router(u), dim=-1))
        regen, regen_report = self.regen(u, retrieved_tokens)
        alpha = self.cfg.residual_alpha or (self.cfg.n_layers ** -0.5)
        fused = torch.tanh(axiom + local + memory + routed + regen)
        hidden = hidden + alpha * fused
        hidden = hidden + alpha * torch.tanh(self.ffn(self.out_norm(hidden)))
        report = {
            "layer_index": self.layer_index,
            "local_exact": local_report,
            "memory": memory_report,
            "regen": regen_report,
            "materializes_effective_dim": False,
            "full_historical_kv_tokens_stored": 0,
        }
        return hidden, next_state, report


class TinyMindAxiomReGenesis(nn.Module):
    """Native TinyMind model: no Mistral/Omega dependency, bounded state, ReGenesis retrieval."""

    def __init__(self, cfg: AxiomReGenesisConfig):
        super().__init__()
        self.cfg = cfg
        self.embed = nn.Embedding(cfg.vocab_size, cfg.dim)
        self.pos = nn.Parameter(torch.zeros(1, cfg.max_seq_len, cfg.dim))
        self.blocks = nn.ModuleList([AxiomReGenesisBlock(cfg, i) for i in range(cfg.n_layers)])
        self.self_assess = SelfAssessmentCore(
            SelfAssessmentCoreConfig(
                dim=cfg.dim,
                inner_dim=cfg.dim * 2,
                recursion_steps=cfg.self_assessment_steps,
                residual_scale=0.10,
                dropout=cfg.dropout,
            )
        )
        self.norm = RMSNorm(cfg.dim)
        self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)
        self.assessment_head = nn.Linear(cfg.dim, 4)

    @property
    def parameter_count(self) -> int:
        return sum(p.numel() for p in self.parameters())

    def parameter_summary(self) -> dict[str, Any]:
        return {
            "architecture": self.cfg.architecture_name,
            "parameter_count": self.parameter_count,
            "layers": self.cfg.n_layers,
            "physical_dim": self.cfg.dim,
            "virtual_dim": self.cfg.axiom_effective_dim,
            "lanes": self.cfg.lanes,
            "materializes_virtual_dim": False,
        }

    def forward(
        self,
        input_ids: torch.Tensor,
        *,
        labels: torch.Tensor | None = None,
        retrieved_tokens: torch.Tensor | None = None,
        states: list[torch.Tensor | None] | None = None,
        return_report: bool = False,
    ) -> dict[str, Any]:
        input_ids = input_ids.clamp(0, self.cfg.vocab_size - 1)
        seq_len = input_ids.shape[1]
        h = self.embed(input_ids) + self.pos[:, :seq_len].to(dtype=self.embed.weight.dtype, device=input_ids.device)
        next_states: list[torch.Tensor] = []
        reports: list[dict[str, Any]] = []
        block_states = states or [None] * len(self.blocks)
        for block, state in zip(self.blocks, block_states):
            h, next_state, report = block(h, state=state, retrieved_tokens=retrieved_tokens)
            next_states.append(next_state)
            if return_report:
                reports.append(report)
        h, assessment = self.self_assess(h)
        logits = self.lm_head(self.norm(h))
        aux = self.assessment_head(h.mean(dim=1))
        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1].contiguous()
            shift_labels = labels[:, 1:].contiguous().clamp(0, self.cfg.vocab_size - 1)
            loss = F.cross_entropy(shift_logits.view(-1, self.cfg.vocab_size), shift_labels.reshape(-1), ignore_index=0)
            evidence_target = torch.ones(aux.shape[0], device=aux.device, dtype=torch.long)
            loss = loss + 0.01 * F.cross_entropy(aux, evidence_target)
        out: dict[str, Any] = {
            "logits": logits,
            "loss": loss,
            "states": next_states,
            "self_assessment": assessment,
            "assessment_logits": aux,
        }
        if return_report:
            out["report"] = {
                "architecture": "TinyMindAxiomReGenesis",
                "config": asdict(self.cfg),
                "parameter_count": self.parameter_count,
                "layer_reports": reports,
                "kv_tokens_stored": 0,
                "materializes_effective_dim": False,
            }
        return out

    def export_runtime_metadata(self, out_dir: str | Path, *, training_report: dict[str, Any] | None = None) -> dict[str, Any]:
        out = Path(out_dir)
        out.mkdir(parents=True, exist_ok=True)
        metadata = {
            "schema": "tinymind.native_axiom_regenesis.runtime_bundle.v1",
            "created_at": datetime.now(timezone.utc).isoformat(),
            "model": self.parameter_summary(),
            "config": asdict(self.cfg),
            "runtime": {
                "format": "TinyMind native PyTorch checkpoint",
                "gguf_standard_compatible": False,
                "gguf_claim_allowed": False,
                "reason": "AxiomReGenesis uses native recurrent/ReGenesis blocks that are not GGUF transformer layers.",
                "kv_tokens_stored_for_long_context": 0,
                "uses_evidence_ledger_retrieval": True,
                "local_exact_window": self.cfg.local_window,
            },
            "training_report_attached": training_report is not None,
            "claim_gate": {
                "native_architecture_independent": True,
                "world_best_claim_allowed": False,
                "external_official_claim_allowed": False,
            },
        }
        if training_report is not None:
            metadata["training_summary"] = training_report.get("summary", {})
            metadata["training_metrics"] = training_report.get("metrics", {})
        path = out / "native_runtime_metadata.json"
        metadata["json_path"] = str(path)
        path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
        return metadata

    @torch.no_grad()
    def generate(self, input_ids: torch.Tensor, *, max_new_tokens: int = 32, retrieved_tokens: torch.Tensor | None = None) -> torch.Tensor:
        self.eval()
        out = input_ids.clone()
        states = None
        for _ in range(max_new_tokens):
            window = out[:, -self.cfg.max_seq_len :]
            result = self(window, retrieved_tokens=retrieved_tokens, states=states)
            states = result["states"]
            next_id = result["logits"][:, -1].argmax(dim=-1, keepdim=True)
            out = torch.cat([out, next_id], dim=1)
        return out


def finite_forward_backward_check(cfg: AxiomReGenesisConfig) -> dict[str, Any]:
    torch.manual_seed(20260528)
    model = TinyMindAxiomReGenesis(cfg)
    ids = torch.randint(4, cfg.vocab_size, (2, min(32, cfg.max_seq_len)))
    retrieved = torch.randint(4, cfg.vocab_size, (2, cfg.regen_top_k, 12))
    out = model(ids, labels=ids, retrieved_tokens=retrieved, return_report=True)
    loss = out["loss"]
    assert loss is not None
    loss.backward()
    grads = [p.grad for p in model.parameters() if p.grad is not None]
    return {
        "forward_finite": bool(torch.isfinite(out["logits"]).all().item()),
        "loss_finite": bool(torch.isfinite(loss).item()),
        "backward_finite": bool(grads) and all(torch.isfinite(g).all().item() for g in grads),
        "parameter_count": model.parameter_count,
        "report": out["report"],
    }


def config_to_dict(cfg: AxiomReGenesisConfig) -> dict[str, Any]:
    return asdict(cfg)


def config_from_dict(payload: dict[str, Any]) -> AxiomReGenesisConfig:
    return AxiomReGenesisConfig(**payload)


AxiomReGenesisModel = TinyMindAxiomReGenesis