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

"""OmniPure data forge plus real TinyMind local train/eval runner."""

from __future__ import annotations

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

from data.omni_pure_forge import OMNI_PURE_DOMAIN_GROUPS, OmniPureForge
from evaluation.knowledge_full_cycle import SourceTraceIndex, audit_pure_records, evaluate_natural_answer_style
from evaluation.local_evidence import run_local_train_eval_bundle


def _read_jsonl(path: str | Path) -> list[dict]:
    p = Path(path)
    return [json.loads(line) for line in p.read_text(encoding="utf-8").splitlines() if line.strip()]


def run_omni_pure_data_train(
    out_dir: str | Path,
    records_per_domain: int = 4,
    train_steps: int = 12,
    seed: int = 20260523,
) -> dict:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    data_dir = out / "dataset"
    train_dir = out / "train_eval"

    manifest = OmniPureForge(records_per_domain=records_per_domain, eval_ratio=0.2).write_jsonl(data_dir)
    rows = _read_jsonl(data_dir / manifest["train_path"]) + _read_jsonl(data_dir / manifest["eval_path"])
    audit = audit_pure_records(rows)
    natural = evaluate_natural_answer_style(rows)
    source_index = SourceTraceIndex.from_records(rows)
    source_meta = source_index.write(out / "omni_pure_source_trace_index.json")
    train = run_local_train_eval_bundle(
        train_dir,
        train_steps=train_steps,
        context_lengths=(32, 128, 1024),
        seed=seed,
        records=rows,
    )

    groups_present = set(manifest["domain_group_counts"])
    coverage_percent = 100.0 * len(groups_present & set(OMNI_PURE_DOMAIN_GROUPS)) / len(OMNI_PURE_DOMAIN_GROUPS)
    report = {
        "schema_version": "tinymind-omni-pure-data-train-report-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "goal": "diverse high-purity multidomain training data with real local train/eval evidence",
        "dataset_manifest": manifest,
        "audit": audit,
        "natural_answer_style": natural,
        "source_trace": {
            **source_meta,
            "example_queries": [
                {"query": "Thai English precise translation", "hits": source_index.query("Thai English precise translation")},
                {"query": "INT6 sparse runtime quantization", "hits": source_index.query("INT6 sparse runtime quantization")},
                {"query": "sandbox Lua Rust tool loop", "hits": source_index.query("sandbox Lua Rust tool loop")},
            ],
        },
        "coverage_gate": {
            "passed": bool(manifest["coverage_gate"]["passed"]),
            "coverage_percent": coverage_percent,
            "domain_groups": sorted(groups_present),
        },
        "unified_representation_space": manifest.get("unified_representation_space", {}),
        "frontier_compression_gate": manifest.get("frontier_compression_gate", {}),
        "pure_gate": {
            "passed": bool(manifest["purity_gate"]["passed"] and audit["blocked_records"] == 0),
            "purity_score": audit["purity_score"],
            "blocked_records": audit["blocked_records"],
        },
        "natural_gate": {"passed": natural["score"] >= 0.75, "score": natural["score"]},
        "train_eval": train.get("train_eval", {}),
        "artifacts": train.get("artifacts", {}),
        "omni_pure_gate": {
            "passed": bool(
                manifest["coverage_gate"]["passed"]
                and manifest.get("unified_representation_space", {}).get("passed", False)
                and manifest.get("frontier_compression_gate", {}).get("passed", False)
                and manifest["purity_gate"]["passed"]
                and audit["blocked_records"] == 0
                and natural["score"] >= 0.75
                and train.get("artifacts")
            ),
            "world_best_claim_allowed": False,
            "definition": (
                "Passed means all declared domain groups and Thai/English coverage are present, all rows pass strict "
                "CEV purity gates, source trace exists, and a real local train/eval bundle was produced."
            ),
        },
    }
    json_path = out / "omni_pure_training_report.json"
    md_path = out / "omni_pure_training_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, sort_keys=True), encoding="utf-8")
    md_path.write_text(_markdown(report), encoding="utf-8")
    return report


def _markdown(report: dict) -> str:
    train = report.get("train_eval", {})
    return "\n".join(
        [
            "# TinyMind OmniPure Data Training",
            "",
            f"- OmniPure gate: {report['omni_pure_gate']['passed']}",
            f"- Pure gate: {report['pure_gate']['passed']} ({report['pure_gate']['purity_score']:.2%})",
            f"- Coverage gate: {report['coverage_gate']['passed']} ({report['coverage_gate']['coverage_percent']:.1f}%)",
            f"- Natural answer gate: {report['natural_gate']['passed']} ({report['natural_gate']['score']:.2%})",
            f"- Records: {report['dataset_manifest']['records_written']}",
            f"- Domain groups: {len(report['coverage_gate']['domain_groups'])}",
            f"- Unified representation space: {report.get('unified_representation_space', {}).get('passed', False)}",
            f"- Frontier compression gate: {report.get('frontier_compression_gate', {}).get('passed', False)}",
            f"- Eval loss: {train.get('eval_loss', 'missing')}",
            f"- Perplexity: {train.get('perplexity', 'missing')}",
            "- World-best claim allowed: false",
            "",
            report["omni_pure_gate"]["definition"],
            "",
        ]
    )