tinymind-native-8b-remote-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"
    tokenizer_mode: str = "byte"
    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
    repeat_unlikelihood_weight: float = 0.02
    entropy_floor_weight: float = 0.001
    entropy_floor: float = 1.25
    default_repetition_penalty: float = 1.18
    default_no_repeat_ngram_size: int = 3


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()
            valid = shift_labels >= 0
            shift_labels = torch.where(valid, shift_labels.clamp(0, self.cfg.vocab_size - 1), shift_labels)
            loss = F.cross_entropy(shift_logits.view(-1, self.cfg.vocab_size), shift_labels.reshape(-1), ignore_index=-100)
            if self.cfg.repeat_unlikelihood_weight > 0:
                probs = torch.softmax(shift_logits, dim=-1)
                prev_tokens = input_ids[:, :-1].clamp(0, self.cfg.vocab_size - 1)
                prev_prob = probs.gather(-1, prev_tokens.unsqueeze(-1)).squeeze(-1)
                repeat_penalty = -torch.log1p(-prev_prob.clamp(max=1 - 1e-6))
                if valid.any():
                    loss = loss + self.cfg.repeat_unlikelihood_weight * repeat_penalty[valid].mean()
            if self.cfg.entropy_floor_weight > 0:
                log_probs = torch.log_softmax(shift_logits, dim=-1)
                probs = log_probs.exp()
                entropy = -(probs * log_probs).sum(dim=-1)
                entropy_penalty = F.relu(float(self.cfg.entropy_floor) - entropy).pow(2)
                if valid.any():
                    loss = loss + self.cfg.entropy_floor_weight * entropy_penalty[valid].mean()
            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

    def _constrain_next_logits(
        self,
        logits: torch.Tensor,
        generated: torch.Tensor,
        *,
        prompt_len: int,
        step: int,
        repetition_penalty: float,
        no_repeat_ngram_size: int,
        min_new_tokens: int,
        valid_token_limit: int,
    ) -> torch.Tensor:
        logits = logits.clone()
        floor = torch.finfo(logits.dtype).min
        if valid_token_limit < self.cfg.vocab_size:
            logits[:, valid_token_limit:] = floor
        if self.cfg.vocab_size > 1:
            logits[:, 0] = floor
            logits[:, 1] = floor
        if step < min_new_tokens and self.cfg.vocab_size > 2:
            logits[:, 2] = floor

        if repetition_penalty > 1.0:
            recent = generated[:, max(0, generated.shape[1] - 96) :]
            for batch_idx in range(generated.shape[0]):
                seen = torch.unique(recent[batch_idx])
                seen = seen[(seen >= 0) & (seen < self.cfg.vocab_size)]
                if seen.numel() == 0:
                    continue
                selected = logits[batch_idx, seen]
                logits[batch_idx, seen] = torch.where(
                    selected > 0,
                    selected / repetition_penalty,
                    selected * repetition_penalty,
                )

        if no_repeat_ngram_size > 1 and generated.shape[1] >= no_repeat_ngram_size - 1:
            n = int(no_repeat_ngram_size)
            for batch_idx in range(generated.shape[0]):
                tokens = generated[batch_idx].detach().cpu().tolist()
                prefix = tuple(tokens[-(n - 1) :])
                banned: set[int] = set()
                for idx in range(0, len(tokens) - n + 1):
                    if tuple(tokens[idx : idx + n - 1]) == prefix:
                        banned.add(int(tokens[idx + n - 1]))
                if banned:
                    banned_tensor = torch.tensor(
                        [token for token in banned if 0 <= token < self.cfg.vocab_size],
                        device=logits.device,
                        dtype=torch.long,
                    )
                    if banned_tensor.numel() > 0:
                        logits[batch_idx, banned_tensor] = floor

        if generated.shape[1] - prompt_len >= 4:
            tail = generated[:, -4:]
            repeated = (tail == tail[:, -1:]).all(dim=1)
            for batch_idx, is_repeated in enumerate(repeated.tolist()):
                if is_repeated:
                    logits[batch_idx, int(tail[batch_idx, -1].item())] = floor

        all_blocked = ~torch.isfinite(logits).any(dim=-1)
        if all_blocked.any():
            logits[all_blocked] = 0
            logits[all_blocked, :valid_token_limit] = 0
            logits[all_blocked, 0:2] = floor
        return logits

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        *,
        max_new_tokens: int = 32,
        retrieved_tokens: torch.Tensor | None = None,
        repetition_penalty: float | None = None,
        no_repeat_ngram_size: int | None = None,
        min_new_tokens: int = 1,
    ) -> torch.Tensor:
        self.eval()
        out = input_ids.clone()
        states = None
        valid_token_limit = min(self.cfg.vocab_size, 260)
        prompt_len = int(input_ids.shape[1])
        penalty = float(repetition_penalty or self.cfg.default_repetition_penalty)
        ngram = int(no_repeat_ngram_size or self.cfg.default_no_repeat_ngram_size)
        for step in range(max_new_tokens):
            window = out[:, -self.cfg.max_seq_len :]
            result = self(window, retrieved_tokens=retrieved_tokens, states=states)
            states = result["states"]
            logits = result["logits"][:, -1].clone()
            logits = self._constrain_next_logits(
                logits,
                out,
                prompt_len=prompt_len,
                step=step,
                repetition_penalty=penalty,
                no_repeat_ngram_size=ngram,
                min_new_tokens=min_new_tokens,
                valid_token_limit=valid_token_limit,
            )
            next_id = logits.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