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

from __future__ import annotations

import argparse
import json
from pathlib import Path
import sys

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

ROOT = Path(r"D:\ad\tinymind\model\tinymind-12b")
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from tinymind_text_sanitize import sanitize_generated_text


PROJECT_ROOT = ROOT.parents[1]


def _default_adapter() -> str:
    interaction_path = PROJECT_ROOT / "reports" / "interaction_champion.json"
    if interaction_path.exists():
        try:
            champion = json.loads(interaction_path.read_text(encoding="utf-8")).get("champion") or {}
            adapter = Path(str(champion.get("adapter") or ""))
            if (adapter / "adapter_config.json").exists():
                return str(adapter)
        except (OSError, json.JSONDecodeError):
            pass
    champion_path = PROJECT_ROOT / "reports" / "system_auto_tuner" / "champion_adapter.json"
    if champion_path.exists():
        try:
            champion = json.loads(champion_path.read_text(encoding="utf-8")).get("champion") or {}
            adapter = Path(str(champion.get("adapter") or ""))
            if (adapter / "adapter_config.json").exists():
                return str(adapter)
        except (OSError, json.JSONDecodeError):
            pass
    return str(ROOT / "adapters" / "tinymind-12b-lora")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("prompt")
    parser.add_argument("--model-id", default="mistralai/Mistral-Nemo-Instruct-2407")
    parser.add_argument("--adapter", default=_default_adapter())
    parser.add_argument("--max-new-tokens", type=int, default=512)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--top-p", type=float, default=0.9)
    parser.add_argument("--repetition-penalty", type=float, default=1.16)
    parser.add_argument("--no-repeat-ngram-size", type=int, default=5)
    args = parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.adapter if Path(args.adapter).exists() else args.model_id, trust_remote_code=True)
    bnb = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        bnb_4bit_use_double_quant=True,
    )
    base = AutoModelForCausalLM.from_pretrained(
        args.model_id,
        quantization_config=bnb,
        device_map="auto",
        trust_remote_code=True,
        dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    )
    model = PeftModel.from_pretrained(base, args.adapter) if Path(args.adapter).exists() else base
    model.eval()

    messages = [
        {"role": "system", "content": "You are TinyMind 12B, precise, tool-aware, evidence-first, and safety-bound."},
        {"role": "user", "content": args.prompt},
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    with torch.no_grad():
        do_sample = args.temperature > 0
        generation_kwargs = {
            "max_new_tokens": args.max_new_tokens,
            "do_sample": do_sample,
            "repetition_penalty": args.repetition_penalty,
            "no_repeat_ngram_size": args.no_repeat_ngram_size,
            "pad_token_id": tokenizer.eos_token_id,
            "eos_token_id": tokenizer.eos_token_id,
        }
        if do_sample:
            generation_kwargs["temperature"] = args.temperature
            generation_kwargs["top_p"] = args.top_p
        output = model.generate(**inputs, **generation_kwargs)
    response = tokenizer.decode(output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    print(sanitize_generated_text(response))
    return 0


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