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