| """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() |
|
|
