"""Run the published Transformers model locally or from Hugging Face."""
from __future__ import annotations
import argparse
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default=str(Path(__file__).resolve().parent))
    parser.add_argument("--prompt", default="你好，请介绍一下自己。")
    parser.add_argument("--device", choices=("auto", "cpu", "cuda", "mps"), default="auto")
    parser.add_argument("--max-new-tokens", type=int, default=256)
    parser.add_argument("--max-denoising-steps", type=int, default=256)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--no-thinking", action="store_true")
    parser.add_argument("--stream", action="store_true")
    args = parser.parse_args()
    device = args.device
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        args.model, trust_remote_code=True, dtype="auto", device_map=device,
        attn_implementation="sdpa",
    ).eval()
    inputs = tokenizer.apply_chat_template(
        [{"role": "user", "content": args.prompt}],
        enable_thinking=not args.no_thinking, add_generation_prompt=True,
        return_tensors="pt", return_dict=True,
    ).to(model.device)
    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) if args.stream else None
    with torch.inference_mode():
        output = model.generate(**inputs, max_new_tokens=args.max_new_tokens,
                                max_denoising_steps=args.max_denoising_steps,
                                seed=args.seed, streamer=streamer)
    if streamer is None:
        print(tokenizer.decode(output.sequences[0, inputs.input_ids.shape[1]:], skip_special_tokens=True))


if __name__ == "__main__":
    main()