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

from __future__ import annotations

from datetime import datetime, timezone
import json
from pathlib import Path
from textwrap import dedent

import torch

from model.architecture import OmegaModel
from model.config import OmegaConfig


def _target_config(physical_layers: int) -> OmegaConfig:
    cfg = OmegaConfig(
        vocab_size=65_536,
        dim=4096,
        n_layers=physical_layers,
        n_heads=32,
        head_dim=128,
        ffn_mult=4,
        layer_pattern="P",
        memory_slots=8,
        memory_ranks=96,
        local_window=2048,
        timescale_count=8,
        low_rank=32,
        residual_alpha=physical_layers ** -0.5,
        max_seq_len=131_072,
        rope_theta=4_000_000.0,
        dropout=0.03,
    )
    cfg.architecture_mode = "purefield"
    cfg.cnn_core_enabled = True
    cfg.self_assessment_enabled = True
    cfg.self_assessment_frequency = 4
    cfg.self_assessment_steps = 3
    cfg.regen_kv_enabled = True
    cfg.max_persistent_tokens = 20_000_000
    cfg.retrieval_top_k = 8
    return cfg


def _smoke_config(smoke_layers: int, smoke_dim: int, smoke_seq_len: int) -> OmegaConfig:
    heads = max(1, min(8, smoke_dim // 64))
    head_dim = smoke_dim // heads
    cfg = OmegaConfig(
        vocab_size=512,
        dim=smoke_dim,
        n_layers=smoke_layers,
        n_heads=heads,
        head_dim=head_dim,
        ffn_mult=2,
        layer_pattern="P",
        memory_ranks=max(8, smoke_dim // 8),
        local_window=min(32, smoke_seq_len),
        timescale_count=4,
        low_rank=4,
        residual_alpha=smoke_layers ** -0.5,
        max_seq_len=max(64, smoke_seq_len),
        dropout=0.0,
    )
    cfg.architecture_mode = "purefield"
    cfg.cnn_core_enabled = True
    cfg.self_assessment_enabled = True
    cfg.self_assessment_frequency = max(1, smoke_layers)
    cfg.self_assessment_steps = 2
    cfg.regen_kv_enabled = False
    return cfg


def _estimate_params(cfg: OmegaConfig) -> int:
    emb = cfg.vocab_size * cfg.dim
    ffn = 2 * cfg.dim * cfg.dim * cfg.ffn_mult
    purefield_shared = (
        cfg.dim * cfg.memory_ranks * 2
        + cfg.dim * cfg.dim
        + cfg.dim * cfg.timescale_count * 2
        + cfg.memory_ranks * cfg.timescale_count
        + (cfg.dim * 3) * cfg.dim
    )
    adapters = cfg.n_layers * (
        2 * cfg.low_rank * (cfg.dim + cfg.memory_ranks)
        + cfg.low_rank * (cfg.dim + cfg.dim)
        + 2 * cfg.low_rank * (cfg.dim + cfg.timescale_count)
        + cfg.low_rank * (cfg.memory_ranks + cfg.timescale_count)
        + cfg.low_rank * (cfg.dim * 3 + cfg.dim)
    )
    per_layer_norms = cfg.n_layers * cfg.dim * 4
    cnn = cfg.dim * cfg.dim * len(cfg.cnn_kernel_sizes) if cfg.cnn_core_enabled else 0
    self_assess = cfg.dim * cfg.dim * cfg.self_assessment_inner_mult * 4 if cfg.self_assessment_enabled else 0
    return int(emb + purefield_shared + adapters + ffn * 0 + per_layer_norms + cnn + self_assess)


def _run_smoke(cfg: OmegaConfig, seq_len: int) -> dict:
    torch.manual_seed(20260527)
    model = OmegaModel(cfg)
    model.train()
    input_ids = torch.randint(0, cfg.vocab_size, (2, seq_len), dtype=torch.long)
    labels = input_ids.clone()
    out = model(input_ids, labels=labels)
    loss = out["loss"]
    forward_finite = bool(torch.isfinite(loss).item())
    loss.backward()
    grad_values = [p.grad.detach().float().abs().mean() for p in model.parameters() if p.grad is not None]
    grad_mean = torch.stack(grad_values).mean() if grad_values else torch.tensor(float("nan"))
    return {
        "forward_finite": forward_finite,
        "backward_finite": bool(torch.isfinite(grad_mean).item()),
        "loss": float(loss.detach().cpu()),
        "grad_abs_mean": float(grad_mean.detach().cpu()),
        "smoke_layers": cfg.n_layers,
        "smoke_dim": cfg.dim,
        "smoke_seq_len": seq_len,
        "smoke_params": sum(p.numel() for p in model.parameters()),
    }


def _write_ffi_artifacts(out: Path, physical_layers: int, virtual_lanes: int) -> dict:
    ffi_dir = out / "ffi"
    ffi_dir.mkdir(parents=True, exist_ok=True)
    virtual_depth = physical_layers * virtual_lanes

    cargo = dedent(
        """\
        [package]
        name = "tinymind_deepweave_t0_ffi"
        version = "0.1.0"
        edition = "2021"

        [lib]
        crate-type = ["cdylib", "rlib"]
        """
    )
    lib_rs = dedent(
        f"""\
        #[no_mangle]
        pub extern "C" fn deepweave_t0_virtual_depth() -> u32 {{
            {virtual_depth}
        }}

        #[no_mangle]
        pub extern "C" fn deepweave_t0_score(raw_score: f32, repair_score: f32) -> f32 {{
            if !raw_score.is_finite() || !repair_score.is_finite() {{
                return 0.0;
            }}
            0.70 * raw_score + 0.30 * repair_score
        }}
        """
    )
    header = dedent(
        """\
        #pragma once
        #include <stdint.h>

        #ifdef __cplusplus
        extern "C" {
        #endif

        uint32_t deepweave_t0_virtual_depth(void);
        float deepweave_t0_score(float raw_score, float repair_score);

        #ifdef __cplusplus
        }
        #endif
        """
    )
    py = dedent(
        """\
        from __future__ import annotations

        import ctypes
        from pathlib import Path


        def load_deepweave_t0(path: str | Path):
            lib = ctypes.CDLL(str(path))
            lib.deepweave_t0_virtual_depth.restype = ctypes.c_uint32
            lib.deepweave_t0_score.argtypes = [ctypes.c_float, ctypes.c_float]
            lib.deepweave_t0_score.restype = ctypes.c_float
            return lib
        """
    )

    (ffi_dir / "Cargo.toml").write_text(cargo, encoding="utf-8")
    (ffi_dir / "src").mkdir(exist_ok=True)
    (ffi_dir / "src" / "lib.rs").write_text(lib_rs, encoding="utf-8")
    (ffi_dir / "tinymind_deepweave_t0.h").write_text(header, encoding="utf-8")
    (ffi_dir / "deepweave_t0_ctypes.py").write_text(py, encoding="utf-8")
    return {
        "abi_name": "tinymind_deepweave_t0_ffi",
        "languages": ["python", "rust", "c_abi"],
        "exports": ["deepweave_t0_virtual_depth", "deepweave_t0_score"],
        "artifacts": {
            "cargo": "ffi/Cargo.toml",
            "rust": "ffi/src/lib.rs",
            "c_header": "ffi/tinymind_deepweave_t0.h",
            "python_ctypes": "ffi/deepweave_t0_ctypes.py",
        },
    }


def _write_markdown(report: dict, path: Path) -> None:
    target = report["target_spec"]
    smoke = report["smoke_evidence"]
    lines = [
        "# TinyMind DeepWeave-T0 Candidate",
        "",
        f"- Physical layers: {target['physical_layers']}",
        f"- Virtual lanes: {target['virtual_lanes']}",
        f"- Virtual tensor depth: {target['virtual_tensor_depth']}",
        f"- Smoke forward finite: {smoke['forward_finite']}",
        f"- Smoke backward finite: {smoke['backward_finite']}",
        f"- Tier-0 claim allowed: {report['claim_gate']['tier0_claim_allowed']}",
        "",
        "This is a native architecture candidate report, not a trained tier-0 model claim.",
    ]
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def build_deepweave_t0_candidate_report(
    out_dir: str | Path,
    *,
    physical_layers: int = 96,
    virtual_lanes: int = 64,
    smoke_layers: int = 4,
    smoke_dim: int = 256,
    smoke_seq_len: int = 16,
) -> dict:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    target_cfg = _target_config(physical_layers)
    smoke_cfg = _smoke_config(smoke_layers, smoke_dim, smoke_seq_len)
    smoke = _run_smoke(smoke_cfg, smoke_seq_len)
    ffi = _write_ffi_artifacts(out, physical_layers, virtual_lanes)

    report = {
        "schema": "tinymind.deepweave_t0_candidate.v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "target_spec": {
            "native_architecture": "TinyMind-DeepWeave-T0-Candidate",
            "physical_layers": physical_layers,
            "virtual_lanes": virtual_lanes,
            "virtual_tensor_depth": physical_layers * virtual_lanes,
            "hidden_dim": target_cfg.dim,
            "heads": target_cfg.n_heads,
            "layer_pattern": target_cfg.layer_pattern,
            "purefield_enabled": target_cfg.architecture_mode == "purefield",
            "cnn_core_enabled": target_cfg.cnn_core_enabled,
            "self_assessment_enabled": target_cfg.self_assessment_enabled,
            "self_assessment_frequency": target_cfg.self_assessment_frequency,
            "regen_kv_enabled": target_cfg.regen_kv_enabled,
            "max_persistent_tokens": target_cfg.max_persistent_tokens,
            "estimated_native_params": _estimate_params(target_cfg),
        },
        "smoke_evidence": smoke,
        "ffi_bridge": ffi,
        "claim_gate": {
            "tier0_claim_allowed": False,
            "world_best_claim_allowed": False,
            "physical_6144_layer_claim_allowed": False,
            "reason": "6144 is virtual tensor depth. This report proves candidate construction and smoke finiteness only.",
        },
    }
    json_path = out / "deepweave_t0_candidate_report.json"
    md_path = out / "deepweave_t0_candidate_report.md"
    report["json_path"] = str(json_path)
    report["markdown_path"] = str(md_path)
    json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    _write_markdown(report, md_path)
    return report