tinymind-native-colab-handoff/bundle/model/tinymind-fusion/compile_fusion_runtime.py

from __future__ import annotations

import argparse
import json
from pathlib import Path

import joblib
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.neighbors import NearestNeighbors


ROOT = Path(r"D:\ad\tinymind\model\tinymind-fusion")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", default=str(ROOT))
    args = parser.parse_args()
    root = Path(args.root)
    src = root / "jsonl" / "fusiondistill_gold.jsonl"
    records = [json.loads(line) for line in src.read_text(encoding="utf-8").splitlines() if line.strip()]
    texts = [f"{r['domain']}\n{r['task']}\n{r['fused_answer']}" for r in records]
    vectorizer = TfidfVectorizer(analyzer="char_wb", ngram_range=(3, 6), max_features=100_000)
    matrix = vectorizer.fit_transform(texts)
    nn = NearestNeighbors(n_neighbors=min(5, len(records)), metric="cosine").fit(matrix)
    artifacts = root / "artifacts"
    artifacts.mkdir(parents=True, exist_ok=True)
    joblib.dump({"vectorizer": vectorizer, "nn": nn, "matrix": matrix, "records": records}, artifacts / "tinymind_fusion_runtime.joblib")
    manifest = {"records": len(records), "artifact": str(artifacts / "tinymind_fusion_runtime.joblib")}
    (artifacts / "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())