File size: 2,112 Bytes
4bbff32 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | """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()
|