tinymind-native-colab-handoff/bundle/model/tinymind-12b/prepare_sft_dataset.py

from __future__ import annotations

import json
from pathlib import Path


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


SYSTEM = (
    "You are TinyMind 12B, a precise tool-aware AI. "
    "Use evidence before speculation. Prefer read-only and reversible actions. "
    "Refuse leaked/private/destructive/license-risk requests and offer lawful alternatives."
)


def read_jsonl(path: Path):
    if not path.exists():
        return
    for line in path.read_text(encoding="utf-8").splitlines():
        if line.strip():
            yield json.loads(line)


def write_record(f, user: str, assistant: str, source: str):
    item = {
        "messages": [
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant},
        ],
        "source": source,
    }
    f.write(json.dumps(item, ensure_ascii=False) + "\n")


def main() -> int:
    out_dir = ROOT / "data"
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / "tinymind_sft.jsonl"
    count = 0

    apex = DATA_ROOT / "distill" / "jsonl" / "apexdistill_gold_10000d.jsonl"
    toolcall = DATA_ROOT / "toolcall" / "jsonl" / "toolcall_gold.jsonl"

    with out_path.open("w", encoding="utf-8") as f:
        for item in read_jsonl(apex) or []:
            user = item.get("task", "")
            synthesis = item.get("synthesis", {})
            answer = synthesis.get("final_answer", "")
            calls = synthesis.get("recommended_tool_calls", [])
            if calls:
                answer += "\n\nSuggested tool calls:\n```json\n" + json.dumps(calls, ensure_ascii=False, indent=2) + "\n```"
            write_record(f, user, answer, "apexdistill_10000d")
            count += 1

        for item in read_jsonl(toolcall) or []:
            messages = item.get("messages", [])
            user = next((m.get("content", "") for m in messages if m.get("role") == "user"), "")
            answer = "Use this structured tool call:\n```json\n" + json.dumps(item.get("expected_tool_calls", []), ensure_ascii=False, indent=2) + "\n```"
            validation = item.get("validation", {})
            if validation:
                answer += "\n\nValidation:\n```json\n" + json.dumps(validation, ensure_ascii=False, indent=2) + "\n```"
            write_record(f, user, answer, "toolcall_gold")
            count += 1

    manifest = {
        "records": count,
        "path": str(out_path),
        "sources": [str(apex), str(toolcall)],
    }
    manifest_dir = ROOT / "manifests"
    manifest_dir.mkdir(parents=True, exist_ok=True)
    (manifest_dir / "sft_dataset_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())