初始化项目,由ModelHub XC社区提供模型
Model: NaClNeZn/Qwen3-1.7B-QLoRA-Shizuku-v1 Source: Original Platform
This commit is contained in:
40
run.py
Normal file
40
run.py
Normal file
@@ -0,0 +1,40 @@
|
||||
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
|
||||
)
|
||||
Reference in New Issue
Block a user