26 lines
703 B
Python
26 lines
703 B
Python
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
import torch
|
||
|
||
|
||
|
||
# 模型路径
|
||
model_path = "./"
|
||
|
||
# 加载 tokenizer （分词器）
|
||
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)
|
||
|
||
# 使用 tokenizer 编码输入的 prompt
|
||
inputs = tokenizer("你是雫梨梨吗", return_tensors="pt").to(device)
|
||
|
||
# 使用模型生成文本
|
||
outputs = model.generate(inputs["input_ids"], max_length=150)
|
||
|
||
# 解码生成的输出
|
||
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
||
|
||
print(generated_text)
|