40 lines
1.0 KiB
Python
40 lines
1.0 KiB
Python
|
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
# 设置模型路径
|
|||
|
|
model_path = "./"
|
|||
|
|
|
|||
|
|
# 加载分词器和模型
|
|||
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
|||
|
|
|
|||
|
|
# 加载模型并移动到可用设备（GPU/CPU）
|
|||
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|||
|
|
model = AutoModelForCausalLM.from_pretrained(model_path).to(device)
|
|||
|
|
|
|||
|
|
# model = AutoModelForCausalLM.from_pretrained(
|
|||
|
|
# model_path,
|
|||
|
|
# trust_remote_code=True,
|
|||
|
|
# torch_dtype=torch.float16,
|
|||
|
|
# device_map="auto"
|
|||
|
|
# )
|
|||
|
|
|
|||
|
|
# 创建流式输出器
|
|||
|
|
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
|||
|
|
|
|||
|
|
# 推理示例
|
|||
|
|
prompt = "喜欢主人吗"
|
|||
|
|
inputs = tokenizer(prompt, return_tensors="pt").to(device)
|
|||
|
|
|
|||
|
|
# 生成（使用 streamer）
|
|||
|
|
print("输入:", prompt)
|
|||
|
|
print("输出: ", end="", flush=True)
|
|||
|
|
|
|||
|
|
outputs = model.generate(
|
|||
|
|
**inputs,
|
|||
|
|
max_new_tokens=512,
|
|||
|
|
# temperature=0.7,
|
|||
|
|
# top_p=0.9,
|
|||
|
|
# repetition_penalty=1.0,
|
|||
|
|
# do_sample=True,
|
|||
|
|
streamer=streamer
|
|||
|
|
)
|