Python加载本地大模型
实现过程
- 加载模型
↓- 加载Tokenizer
↓- 构造Prompt
↓- 文本转Token
↓- GPU推理(generate)
↓- Token转文本(decode)
流程说明:
- 加载模型 :从本地目录读取预训练好的大模型权重文件(如
model.safetensors),并加载到内存中。 - 加载 Tokenizer:分词器负责将自然语言文本与模型能理解的 Token 编号相互转换,是文本处理的第一步。
- 构造 Prompt :将用户输入按照模型要求的聊天模板(如 Qwen 的
<|im_start|>格式)拼接成完整的提示词。 - 文本转 Token:通过 Tokenizer 将 Prompt 文本切分为 Token ID 序列,并转换为 PyTorch Tensor,供模型读取。
- GPU 推理(generate):将 Token 序列送入模型,在 GPU 上执行前向计算,逐 Token 生成回答内容。
- Token 转文本(decode):将模型输出的 Token ID 序列解码回可读的中文文本,得到最终回答。## 所需依赖
javascript
pip install -i https://mirror.nju.edu.cn/pypi/web/simple transformers
pip install -i https://mirror.nju.edu.cn/pypi/web/simple accelerate
代码实现
javascript
import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
# =====================================================
# 指定使用哪块GPU
# 0表示第一块显卡
# =====================================================
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# =====================================================
# 本地模型目录
# 例如:
# /hy-tmp/
# ├── config.json
# ├── tokenizer.json
# ├── model.safetensors
# =====================================================
model_path = "/hy-tmp"
# =====================================================
# 加载Tokenizer(分词器)
#
# 作用:
# 把中文转换成模型认识的Token编号
#
# 示例:
# "你好"
# ↓
# [151644, 872]
# =====================================================
tokenizer = AutoTokenizer.from_pretrained(model_path)
# =====================================================
# 加载大模型
#
# torch_dtype=float16
# 使用FP16半精度
# 显存占用更小
#
# device_map="auto"
# 自动加载到GPU
# 单卡、多卡都能适配
# =====================================================
model = AutoModelForCausalLM.from_pretrained(
model_path,
torch_dtype=torch.float16,
device_map="cuda:0"
)
# =====================================================
# 查看模型实际运行设备
# 正常应输出:
# cuda:0
# =====================================================
print("模型所在设备:", next(model.parameters()).device)
# =====================================================
# 用户问题
# =====================================================
query = "你是谁"
# =====================================================
# 构造聊天格式
#
# role:
# user 用户
# assistant AI
# system 系统提示词
# =====================================================
messages = [
{
"role": "user",
"content": query
}
]
# =====================================================
# 应用Chat模板
#
# Qwen要求:
#
# <|im_start|>user
# 你是谁
# <|im_end|>
#
# 自动拼接成模型需要的格式
# =====================================================
prompt = tokenizer.apply_chat_template(
messages,
tokenize=False, # 返回文本
add_generation_prompt=True # 自动加assistant开始标记
)
print("Prompt内容:")
print(prompt)
# =====================================================
# 分词(Tokenization)
#
# 文本
# ↓
# Token ID
# ↓
# Tensor
#
# return_tensors="pt"
# 表示返回PyTorch Tensor
# =====================================================
inputs = tokenizer(
prompt,
return_tensors="pt"
)
# =====================================================
# 把输入移动到模型所在设备
#
# CPU -> GPU
# =====================================================
inputs = inputs.to(model.device)
# =====================================================
# 模型推理
#
# max_new_tokens
# 最多生成100个Token
#
# do_sample=True
# 开启随机采样
#
# temperature=0.7
# 控制创造性
#
# 0.1 非常保守
# 0.7 常用
# 1.2 发散
# =====================================================
outputs = model.generate(
**inputs,
max_new_tokens=100,
do_sample=True,
temperature=0.7
)
# =====================================================
# 解码
#
# Token
# ↓
# 中文文本
#
# skip_special_tokens=True
# 去掉特殊标记
# =====================================================
answer = tokenizer.decode(
outputs[0],
skip_special_tokens=True
)
print("\n模型回答:")
print(answer)