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

"""Evidence dossier for the AxiomWeave architecture."""

from __future__ import annotations

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

import torch

from model.axiom_weave import AxiomWeaveModel
from model.config import axiomweave_config


def build_axiomweave_dossier(out_dir: str | Path, size: str = "tiny", seq_len: int = 16) -> dict:
    cfg = axiomweave_config(size)
    model = AxiomWeaveModel(cfg)
    input_ids = torch.randint(4, min(cfg.vocab_size, 2048), (1, int(seq_len)))
    with torch.no_grad():
        out = model(input_ids, return_stats=True)
    logits = out["logits"]
    route_means = []
    route_entropy = []
    for row in out["stats"]:
        route_means.append([float(x) for x in row["route_weights_mean"].cpu()])
        route_entropy.append(float(row["route_entropy"].cpu()))
    dossier = {
        "schema_version": "tinymind-axiomweave-dossier-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "architecture": "AxiomWeave",
        "size": size,
        "params": model.count_params_int(),
        "forward_finite": bool(torch.isfinite(logits).all().item()),
        "logits_shape": list(logits.shape),
        "synthesis_axes": [
            "local_attention_for_recent_detail",
            "selective_state_space_for_dynamical_compression",
            "purefield_for_bounded_long_memory",
            "kan_ffn_for_parameter_efficient_nonlinearity",
            "regenesis_ready_exact_archive_for_10m_context",
            "logic_tool_grounding_compatibility",
            "int4_sparse_readiness",
        ],
        "route_weights_mean_by_layer": route_means,
        "route_entropy_mean": sum(route_entropy) / max(len(route_entropy), 1),
        "claim_gate": {
            "new_architecture_in_repo": True,
            "world_first_claim_allowed": False,
            "world_best_claim_allowed": False,
            "reason": "Architecture smoke is local evidence only; external benchmark rank and ablations are required.",
        },
    }
    out_path = Path(out_dir)
    out_path.mkdir(parents=True, exist_ok=True)
    json_path = out_path / "axiomweave_dossier.json"
    md_path = out_path / "axiomweave_dossier.md"
    dossier["json_path"] = str(json_path)
    dossier["markdown_path"] = str(md_path)
    json_path.write_text(json.dumps(dossier, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    md_path.write_text(_markdown(dossier), encoding="utf-8")
    return dossier


def _markdown(dossier: dict) -> str:
    lines = [
        "# TinyMind AxiomWeave Dossier",
        "",
        f"- Params: {dossier['params']:,}",
        f"- Forward finite: {dossier['forward_finite']}",
        f"- Route entropy mean: {dossier['route_entropy_mean']:.4f}",
        "- World-best claim: blocked until external evidence",
        "",
        "## Synthesis Axes",
        "",
    ]
    for axis in dossier["synthesis_axes"]:
        lines.append(f"- {axis}")
    return "\n".join(lines) + "\n"