"""Run the local text-only model: python inference.py '你好' --device cpu"""
import argparse
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('prompt', nargs='?', default='请用一句话介绍你自己。')
    parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
    parser.add_argument('--max-new-tokens', type=int, default=128)
    args = parser.parse_args()
    torch.set_num_threads(8)
    path = Path(__file__).resolve().parent
    tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=True)
    model = AutoModelForCausalLM.from_pretrained(
        path, trust_remote_code=True, local_files_only=True, dtype=torch.bfloat16,
    ).to(args.device).eval()
    text = tokenizer.apply_chat_template(
        [{'role': 'user', 'content': args.prompt}], tokenize=False,
        add_generation_prompt=True, enable_thinking=False,
    )
    inputs = tokenizer(text, return_tensors='pt', add_special_tokens=False).to(args.device)
    inputs.pop('token_type_ids', None)
    with torch.inference_mode():
        tokens = model.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False)
    print(tokenizer.decode(tokens[0, inputs.input_ids.shape[1]:], skip_special_tokens=True))


if __name__ == '__main__':
    main()