Ubuntu下使用Transformers推理框架进行LLM的推理
前期准备：
Nvidia驱动已经安装完成
（nvidia-smi）
CUDA已经安装完成
(nvcc --version)
---------------------------------
1、创建并激活虚拟环境：
conda create -n transformers python=3.10 -y && conda activate transformers
2、部署transformers推理框架：
pip install "transformers[torch]" -i https://pypi.tuna.tsinghua.edu.cn/simple
(未翻墙则加上清华源，如果已科学上网则无需加)
使用pip show transformers命令来检查是否安装成功以及版本信息。
3、安装huggingface下载依赖：
pip install -U huggingface_hub -i https://pypi.tuna.tsinghua.edu.cn/simple
(未翻墙则加上清华源，如果已科学上网则无需加)
4、设置huggingface镜像站的环境变量：
export HF_ENDPOINT=https://hf-mirror.com
（如果已科学上网，则无需此步骤）
5、编写app.py代码，如下：
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
def main():
# 模型名称，可以替换为本地路径，例如 "./Qwen2.5-1.5B"
model_name = "Qwen/Qwen3-1.7B"
# 加载 tokenizer 和模型
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16, # 如果显卡支持半精度，节省显存
device_map="auto", # 自动选择设备（GPU / CPU）
trust_remote_code=True
)
# 输入问题
question = "如何让打糕做的更美味？"
# 编码输入
inputs = tokenizer(question, return_tensors="pt").to(model.device)
# 推理生成
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=300, # 生成的最大字数
temperature=0.7, # 控制生成多样性
top_p=0.9, # nucleus sampling
do_sample=True
)
# 解码输出
answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("模型回答：\n", answer)
if __name__ == "__main__":
main()
6、运行python代码：
python app.py
--------------------------------
流式推理的代码：
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
import torch
import threading
def main():
# 模型路径（可以换成本地路径）
model_name = "Qwen/Qwen3-1.7B"
# 加载 tokenizer 和 模型
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16, # 半精度减少显存
device_map="auto",
trust_remote_code=True
)
# 输入问题
question = "如何让打糕做的更美味？"
# 编码输入
inputs = tokenizer(question, return_tensors="pt").to(model.device)
# 设置流式输出
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
# 推理参数
generation_kwargs = dict(
**inputs,
max_new_tokens=32000,
temperature=0.7,
top_p=0.9,
do_sample=True,
streamer=streamer
)
# 用线程异步生成
thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
print("模型回答：", end="", flush=True)
for new_text in streamer:
print(new_text, end="", flush=True)
thread.join()
if __name__ == "__main__":
main()
评论
发表评论