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

from __future__ import annotations

from datetime import datetime, timezone
import json
from pathlib import Path
import shutil
import zipfile
from typing import Any


ROOT = Path(__file__).resolve().parents[1]


def _copy_tree(src: Path, dst: Path, patterns: tuple[str, ...] = ("*.py",)) -> None:
    dst.mkdir(parents=True, exist_ok=True)
    for pattern in patterns:
        for path in src.rglob(pattern):
            if "__pycache__" in path.parts:
                continue
            target = dst / path.relative_to(src)
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(path, target)


def _remote_train_command(dataset_name: str) -> str:
    return (
        "python -m train.cli native-axiom-regenesis-train "
        "--dataset /content/tinymind_remote/bundle/data/"
        f"{dataset_name} "
        "--out-dir /content/tinymind_axiom_regenesis_8b_target "
        "--max-steps 50000 --eval-records 256 --limit-records 10000 "
        "--dim 2816 --layers 48 --lanes 64 --seq-len 1024 --vocab-size 4096 "
        "--tokenizer-mode char_v1 --virtual-dim 1048576 --basis-rank 256 --facets 64 "
        "--learning-rate 1.2e-05 --train-batch-size 1 --device cuda"
    )


def _baseline_probe_command() -> str:
    return (
        "python -m train.cli native-baseline-probe "
        "--out-dir /content/tinymind_axiom_regenesis_8b_probe "
        "--native-checkpoint /content/tinymind_axiom_regenesis_8b_target/checkpoint.pt "
        "--baseline-report /content/tinymind_remote/bundle/baseline/deep_core_probe_report.json "
        "--max-new-tokens 160 --device cuda"
    )


def _broad_probe_command() -> str:
    return (
        "python -m train.cli native-broad-probe "
        "--out-dir /content/tinymind_axiom_regenesis_8b_broad_probe "
        "--native-checkpoint /content/tinymind_axiom_regenesis_8b_target/checkpoint.pt "
        "--max-new-tokens 192 --device cuda"
    )


def build_native_8b_remote_handoff(
    out_dir: str | Path,
    *,
    dataset: str | Path = "reports/omni_round_curriculum_xl_latest/omni_round_curriculum.jsonl",
) -> dict[str, Any]:
    out = Path(out_dir)
    if out.exists():
        shutil.rmtree(out)
    out.mkdir(parents=True, exist_ok=True)
    bundle = out / "bundle"
    bundle.mkdir(parents=True, exist_ok=True)

    for folder in ("model", "train", "evaluation", "data"):
        _copy_tree(ROOT / folder, bundle / folder)
    (bundle / "scripts").mkdir(parents=True, exist_ok=True)
    for script in ("build_omni_round_curriculum.py", "deep_core_probe.py"):
        source = ROOT / "scripts" / script
        if source.exists():
            shutil.copy2(source, bundle / "scripts" / script)
    for file_name in ("conftest.py", "pytest.ini", "requirements.txt"):
        source = ROOT / file_name
        if source.exists():
            shutil.copy2(source, bundle / file_name)

    data_dir = bundle / "data"
    data_dir.mkdir(parents=True, exist_ok=True)
    dataset_src = ROOT / dataset
    dataset_dst = data_dir / dataset_src.name
    shutil.copy2(dataset_src, dataset_dst)

    baseline_dir = bundle / "baseline"
    baseline_dir.mkdir(parents=True, exist_ok=True)
    baseline_src = ROOT / "reports" / "deep_core_probe_deepsharp_command_mythos_20260526_2320" / "deep_core_probe_report.json"
    if baseline_src.exists():
        shutil.copy2(baseline_src, baseline_dir / "deep_core_probe_report.json")

    train_command = _remote_train_command(dataset_dst.name)
    baseline_command = _baseline_probe_command()
    broad_command = _broad_probe_command()

    notebook = {
        "cells": [
            {
                "cell_type": "markdown",
                "metadata": {},
                "source": [
                    "# TinyMind AxiomReGenesis 8B Target Remote Training\n",
                    "This notebook trains the 8.22B-class native target. It does not claim quality until probes pass.\n",
                ],
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "from google.colab import files\n",
                    "uploaded = files.upload()\n",
                    "assert 'tinymind_native_8b_remote_bundle.zip' in uploaded\n",
                    "!rm -rf /content/tinymind_remote && mkdir -p /content/tinymind_remote\n",
                    "!unzip -q tinymind_native_8b_remote_bundle.zip -d /content/tinymind_remote\n",
                    "%cd /content/tinymind_remote/bundle\n",
                ],
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "!pip -q install -r requirements.txt || true\n",
                    "import torch, sys\n",
                    "print('python', sys.version)\n",
                    "print('cuda available', torch.cuda.is_available())\n",
                    "print('gpu', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'cpu')\n",
                ],
            },
            {"cell_type": "code", "execution_count": None, "metadata": {}, "outputs": [], "source": [train_command + "\n"]},
            {"cell_type": "code", "execution_count": None, "metadata": {}, "outputs": [], "source": [baseline_command + "\n"]},
            {"cell_type": "code", "execution_count": None, "metadata": {}, "outputs": [], "source": [broad_command + "\n"]},
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "!zip -qr /content/tinymind_axiom_regenesis_8b_results.zip /content/tinymind_axiom_regenesis_8b_target /content/tinymind_axiom_regenesis_8b_probe /content/tinymind_axiom_regenesis_8b_broad_probe\n",
                    "files.download('/content/tinymind_axiom_regenesis_8b_results.zip')\n",
                ],
            },
        ],
        "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}},
        "nbformat": 4,
        "nbformat_minor": 5,
    }
    notebook_path = out / "native_axiom_regenesis_8b_remote_train.ipynb"
    notebook_path.write_text(json.dumps(notebook, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    hf_notebook = json.loads(json.dumps(notebook))
    hf_notebook["cells"][1]["source"] = [
        "!pip -q install -U 'huggingface_hub[hf_transfer]'\n",
        "!rm -rf /content/tinymind_remote && mkdir -p /content/tinymind_remote\n",
        "!hf sync hf://buckets/bbkdevops/unicosys-hypergraph-bucket/tinymind-native-8b-remote-handoff /content/tinymind_remote/handoff\n",
        "!unzip -q /content/tinymind_remote/handoff/tinymind_native_8b_remote_bundle.zip -d /content/tinymind_remote\n",
        "%cd /content/tinymind_remote/bundle\n",
    ]
    hf_notebook_path = out / "native_axiom_regenesis_8b_remote_train_from_hf_bucket.ipynb"
    hf_notebook_path.write_text(json.dumps(hf_notebook, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    zip_path = out / "tinymind_native_8b_remote_bundle.zip"
    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=6) as archive:
        for path in bundle.rglob("*"):
            if path.is_file():
                archive.write(path, path.relative_to(out))

    run_script = out / "run_remote_8b_target.ps1"
    run_script.write_text(
        "\n".join(
            [
                "# Run inside the extracted remote bundle on a CUDA host.",
                "Set-Location /content/tinymind_remote/bundle",
                train_command,
                baseline_command,
                broad_command,
                "",
            ]
        ),
        encoding="utf-8",
    )

    manifest = {
        "schema": "tinymind.native_8b_remote_handoff.v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "dataset": str(dataset_src),
        "artifacts": {
            "bundle_zip": str(zip_path),
            "notebook": str(notebook_path),
            "hf_bucket_notebook": str(hf_notebook_path),
            "run_script": str(run_script),
        },
        "hf_bucket_uri": "hf://buckets/bbkdevops/unicosys-hypergraph-bucket/tinymind-native-8b-remote-handoff",
        "commands": {
            "train": train_command,
            "baseline_probe": baseline_command,
            "broad_probe": broad_command,
        },
        "target": {
            "estimated_parameters": 8_219_629_572,
            "layers": 48,
            "dim": 2816,
            "lanes": 64,
            "seq_len": 1024,
            "vocab_size": 4096,
            "virtual_dim": 1_048_576,
        },
        "claim_gate": {
            "remote_handoff_ready": zip_path.exists() and notebook_path.exists(),
            "remote_training_completed": False,
            "external_eval_completed": False,
            "quality_above_larger_models_claim_allowed": False,
            "world_best_claim_allowed": False,
        },
    }
    manifest_path = out / "native_8b_remote_handoff_manifest.json"
    manifest["json_path"] = str(manifest_path)
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return manifest