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

from __future__ import annotations

import argparse
import json
import time
import urllib.request
from pathlib import Path
from typing import Any


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


SEED_TASKS = [
    {
        "id": "fusion-windows-readonly-audit",
        "domain": "windows_system",
        "prompt": "ออกแบบขั้นตอนตรวจ RAM, driver, event log และ power แบบ read-only บน Windows พร้อมคำสั่ง PowerShell ที่ปลอดภัย",
    },
    {
        "id": "fusion-toolcall-schema",
        "domain": "tool_calling",
        "prompt": "อธิบายวิธีเลือก function call ที่ถูก schema สำหรับงาน audit JSONL dataset และยกตัวอย่าง JSON tool call",
    },
    {
        "id": "fusion-thai-data-purity",
        "domain": "thai_data",
        "prompt": "ออกแบบ policy รวบรวมข้อมูลภาษาไทยสำหรับ training ที่คุณภาพสูงและไม่ปนข้อมูลหลุดหรือผิดสิทธิ์",
    },
    {
        "id": "fusion-safe-reverse-engineering",
        "domain": "reverse_engineering",
        "prompt": "วาง workflow reverse engineering binary ที่ได้รับอนุญาตโดยไม่ execute บน host และแยก fact/inference/unknown",
    },
    {
        "id": "fusion-code-repair-loop",
        "domain": "software_engineering",
        "prompt": "ออกแบบ loop แก้บั๊กที่อ่านโค้ดก่อน แก้แบบ minimal patch แล้ว verify ทุกครั้ง",
    },
]


SYSTEM = """You are contributing to TinyMind FusionDistill.
Answer only for lawful, authorized, defensive, or engineering use.
Do not provide leaked data, credential theft, stealth, persistence, exploit chains, DRM bypass, or destructive commands.
Prefer evidence, read-only diagnostics, reversible steps, and exact tool-call JSON when useful.
Return concise high-signal knowledge."""


def ollama_tags() -> list[dict[str, Any]]:
    with urllib.request.urlopen("http://127.0.0.1:11434/api/tags", timeout=10) as response:
        return json.loads(response.read().decode("utf-8")).get("models", [])


def choose_models(spec: str, max_models: int | None = None) -> list[str]:
    models = [m["name"] for m in ollama_tags()]
    if spec != "all":
        wanted = [x.strip() for x in spec.split(",") if x.strip()]
        models = [m for m in models if m in wanted or m.split(":")[0] in wanted]
    # Prefer diverse, capable local models first. Keep all when requested with no max.
    priority = [
        "tinymind-apex-local:latest",
        "deepseek-r1:14b",
        "starcoder2:15b",
        "gemma-4:latest",
        "qwen2.5:7b",
        "llama3.1:8b",
        "mixtral:8x7b",
    ]
    ordered = [m for m in priority if m in models] + [m for m in models if m not in priority]
    return ordered[:max_models] if max_models else ordered


def call_model(model: str, prompt: str, timeout: int) -> dict[str, Any]:
    payload = {
        "model": model,
        "prompt": f"{SYSTEM}\n\nTask:\n{prompt}",
        "stream": False,
        "options": {"temperature": 0.25, "top_p": 0.9, "num_ctx": 4096},
    }
    req = urllib.request.Request(
        "http://127.0.0.1:11434/api/generate",
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    started = time.perf_counter()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as response:
            data = json.loads(response.read().decode("utf-8"))
        return {
            "model": model,
            "ok": True,
            "latency_sec": round(time.perf_counter() - started, 3),
            "response": data.get("response", "").strip(),
        }
    except Exception as exc:
        return {
            "model": model,
            "ok": False,
            "latency_sec": round(time.perf_counter() - started, 3),
            "error": str(exc),
            "response": "",
        }


def fuse(task: dict[str, Any], responses: list[dict[str, Any]]) -> dict[str, Any]:
    good = [r for r in responses if r.get("ok") and r.get("response")]
    blocked_terms = ["password", "ข้อมูลหลุด", "exfiltrate", "disable antivirus", "delete system32"]
    combined = "\n\n".join(f"## {r['model']}\n{r['response']}" for r in good)
    safety_flags = sorted({term for term in blocked_terms if term.lower() in combined.lower()})
    lines = []
    for r in good:
        for line in r["response"].splitlines():
            s = line.strip()
            if len(s) > 20 and s not in lines:
                lines.append(s)
            if len(lines) >= 20:
                break
        if len(lines) >= 20:
            break
    fused = "\n".join(f"- {line.lstrip('- ').strip()}" for line in lines[:16])
    if safety_flags:
        fused = "Safety flags detected; use lawful/read-only alternatives only.\n" + fused
    return {
        "id": task["id"],
        "domain": task["domain"],
        "task": task["prompt"],
        "models_used": [r["model"] for r in good],
        "model_failures": [r for r in responses if not r.get("ok")],
        "raw_responses": responses,
        "fused_answer": fused,
        "quality": {
            "response_count": len(good),
            "safety_flags": safety_flags,
            "fusion_method": "multi_model_extract_distinct_lines",
            "ready_for_runtime": len(good) > 0 and not safety_flags,
        },
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--models", default="tinymind-apex-local:latest,deepseek-r1:14b,starcoder2:15b")
    parser.add_argument("--max-models", type=int)
    parser.add_argument("--task-limit", type=int, default=5)
    parser.add_argument("--timeout", type=int, default=120)
    args = parser.parse_args()

    out_dir = ROOT / "jsonl"
    manifest_dir = ROOT / "manifests"
    out_dir.mkdir(parents=True, exist_ok=True)
    manifest_dir.mkdir(parents=True, exist_ok=True)

    models = choose_models(args.models, args.max_models)
    tasks = SEED_TASKS[: args.task_limit]
    fused_records = []
    for task in tasks:
        responses = [call_model(model, task["prompt"], args.timeout) for model in models]
        fused_records.append(fuse(task, responses))

    out_path = out_dir / "fusiondistill_gold.jsonl"
    with out_path.open("w", encoding="utf-8") as f:
        for record in fused_records:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    audit = {
        "records": len(fused_records),
        "models_requested": models,
        "tasks": [t["id"] for t in tasks],
        "jsonl": str(out_path),
        "successful_model_responses": sum(r["quality"]["response_count"] for r in fused_records),
    }
    (manifest_dir / "fusiondistill_audit.json").write_text(json.dumps(audit, indent=2, ensure_ascii=False), encoding="utf-8")
    print(json.dumps(audit, indent=2, ensure_ascii=False))
    return 0


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