tinymind-native-8b-remote-handoff/bundle/model/tinymind-fusion/run_fusion.py

from __future__ import annotations

import argparse
import json
from pathlib import Path

import joblib


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


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("query")
    args = parser.parse_args()
    bundle = joblib.load(ROOT / "artifacts" / "tinymind_fusion_runtime.joblib")
    q = bundle["vectorizer"].transform([args.query])
    distances, indices = bundle["nn"].kneighbors(q)
    matches = []
    answers = []
    for dist, idx in zip(distances[0], indices[0]):
        r = bundle["records"][int(idx)]
        matches.append({"id": r["id"], "domain": r["domain"], "score": round(1 - float(dist), 6), "models_used": r["models_used"]})
        answers.append(r["fused_answer"])
    print(json.dumps({"answer": "\n\n".join(answers[:2]), "matches": matches}, ensure_ascii=False, indent=2))
    return 0


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