tinymind-native-colab-handoff/bundle/model/tinymind-apex/train_neural_core.py

from __future__ import annotations

import argparse
import json
import random
import time
from collections import Counter
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset


ROOT = Path(r"D:\ad\tinymind\model\tinymind-apex")
ARTIFACTS = ROOT / "artifacts"
MAX_LEN = 512


def load_records() -> list[dict[str, Any]]:
    path = ARTIFACTS / "records.jsonl"
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]


def label_for_tool(record: dict[str, Any]) -> str:
    calls = record.get("tool_calls", [])
    if not calls:
        return "none"
    first = calls[0]
    name = first.get("name") or first.get("tool") or "none"
    return name


def build_vocab(records: list[dict[str, Any]], vocab_size: int = 4096) -> dict[str, int]:
    counter = Counter()
    for record in records:
        counter.update(record["text"])
    vocab = {"<pad>": 0, "<unk>": 1}
    for ch, _ in counter.most_common(vocab_size - len(vocab)):
        if ch not in vocab:
            vocab[ch] = len(vocab)
    return vocab


class TinyMindDataset(Dataset):
    def __init__(self, records: list[dict[str, Any]], vocab: dict[str, int], domains: dict[str, int], tools: dict[str, int]):
        self.records = records
        self.vocab = vocab
        self.domains = domains
        self.tools = tools

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, idx: int):
        record = self.records[idx]
        text = record["text"][:MAX_LEN]
        ids = [self.vocab.get(ch, 1) for ch in text]
        if len(ids) < MAX_LEN:
            ids += [0] * (MAX_LEN - len(ids))
        risk = 1 if str(record.get("quality", {}).get("risk", "")).lower() == "high" else 0
        return {
            "x": torch.tensor(ids, dtype=torch.long),
            "domain": torch.tensor(self.domains[record["domain"]], dtype=torch.long),
            "tool": torch.tensor(self.tools[label_for_tool(record)], dtype=torch.long),
            "risk": torch.tensor(risk, dtype=torch.float32),
        }


class NeuralCore(nn.Module):
    def __init__(self, vocab_size: int, domain_count: int, tool_count: int, emb: int = 96, hidden: int = 192):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, emb, padding_idx=0)
        self.conv3 = nn.Conv1d(emb, hidden, kernel_size=3, padding=1)
        self.conv5 = nn.Conv1d(emb, hidden, kernel_size=5, padding=2)
        self.norm = nn.LayerNorm(hidden * 2)
        self.dropout = nn.Dropout(0.15)
        self.domain_head = nn.Linear(hidden * 2, domain_count)
        self.tool_head = nn.Linear(hidden * 2, tool_count)
        self.risk_head = nn.Linear(hidden * 2, 1)

    def encode(self, x):
        emb = self.embedding(x).transpose(1, 2)
        h3 = torch.relu(self.conv3(emb)).amax(dim=2)
        h5 = torch.relu(self.conv5(emb)).amax(dim=2)
        h = torch.cat([h3, h5], dim=1)
        h = self.dropout(self.norm(h))
        return h

    def forward(self, x):
        h = self.encode(x)
        return {
            "domain": self.domain_head(h),
            "tool": self.tool_head(h),
            "risk": self.risk_head(h).squeeze(-1),
        }


def evaluate(model: NeuralCore, loader: DataLoader, device: str) -> dict[str, float]:
    model.eval()
    total = 0
    domain_ok = 0
    tool_ok = 0
    risk_ok = 0
    with torch.no_grad():
        for batch in loader:
            x = batch["x"].to(device)
            domain = batch["domain"].to(device)
            tool = batch["tool"].to(device)
            risk = batch["risk"].to(device)
            out = model(x)
            domain_ok += (out["domain"].argmax(dim=1) == domain).sum().item()
            tool_ok += (out["tool"].argmax(dim=1) == tool).sum().item()
            risk_pred = (torch.sigmoid(out["risk"]) >= 0.5).float()
            risk_ok += (risk_pred == risk).sum().item()
            total += x.shape[0]
    return {
        "domain_acc": round(domain_ok / max(total, 1), 4),
        "tool_acc": round(tool_ok / max(total, 1), 4),
        "risk_acc": round(risk_ok / max(total, 1), 4),
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", default=str(ROOT))
    parser.add_argument("--epochs", type=int, default=8)
    parser.add_argument("--batch-size", type=int, default=64)
    args = parser.parse_args()

    root = Path(args.root)
    artifacts = root / "artifacts"
    records = load_records()
    random.Random(42).shuffle(records)
    vocab = build_vocab(records)
    domains = {name: i for i, name in enumerate(sorted({r["domain"] for r in records}))}
    tools = {name: i for i, name in enumerate(sorted({label_for_tool(r) for r in records}))}

    split = max(1, int(len(records) * 0.85))
    train_records = records[:split]
    val_records = records[split:]
    train_ds = TinyMindDataset(train_records, vocab, domains, tools)
    val_ds = TinyMindDataset(val_records, vocab, domains, tools)
    train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True)
    val_loader = DataLoader(val_ds, batch_size=args.batch_size)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = NeuralCore(len(vocab), len(domains), len(tools)).to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-3)
    domain_loss = nn.CrossEntropyLoss()
    tool_loss = nn.CrossEntropyLoss()
    risk_loss = nn.BCEWithLogitsLoss()

    history = []
    for epoch in range(1, args.epochs + 1):
        model.train()
        total_loss = 0.0
        for batch in train_loader:
            x = batch["x"].to(device)
            domain = batch["domain"].to(device)
            tool = batch["tool"].to(device)
            risk = batch["risk"].to(device)
            out = model(x)
            loss = domain_loss(out["domain"], domain) + tool_loss(out["tool"], tool) + risk_loss(out["risk"], risk)
            opt.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            total_loss += loss.item()
        metrics = evaluate(model, val_loader, device)
        metrics["epoch"] = epoch
        metrics["loss"] = round(total_loss / max(len(train_loader), 1), 4)
        history.append(metrics)
        print(json.dumps(metrics, ensure_ascii=False))

    neural_dir = artifacts / "neural_core"
    neural_dir.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "state_dict": model.cpu().state_dict(),
            "vocab": vocab,
            "domains": domains,
            "tools": tools,
            "max_len": MAX_LEN,
            "config": {"emb": 96, "hidden": 192},
        },
        neural_dir / "tinymind_neural_core.pt",
    )
    manifest = {
        "name": "tinymind-neural-core",
        "type": "pytorch_char_cnn_multihead",
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "records": len(records),
        "device": device,
        "epochs": args.epochs,
        "history": history,
        "weights": str(neural_dir / "tinymind_neural_core.pt"),
    }
    (neural_dir / "manifest.json").write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
    print(json.dumps(manifest, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())