如何在MPT-7B语言模型上执行推理任务？

前置准备 #

• 安装依赖库：执行pip install transformers torch accelerate，如果要启用8位量化节省显存，额外安装bitsandbytes
• 硬件要求：至少13GB显存（FP16精度），用CPU推理也可行但速度极慢

加载模型与Tokenizer #

MPT-7B采用自定义模型架构，加载时必须开启trust_remote_code=True才能加载对应代码：

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 加载tokenizer和模型
tokenizer = AutoTokenizer.from_pretrained("mosaicml/mpt-7b")
model = AutoModelForCausalLM.from_pretrained(
    "mosaicml/mpt-7b",
    trust_remote_code=True,
    torch_dtype=torch.float16,
    device_map="auto"  # 自动分配模型到GPU/CPU
)

显存不足时，可添加load_in_8bit=True参数启用8位量化，能把显存占用压到8GB左右。

执行文本生成推理 #

以下是基础推理代码，输入提示文本后自动生成后续内容：

# 输入提示文本
prompt = "今天天气很好，适合"

# 编码输入并移至模型所在设备
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

# 生成后续内容
outputs = model.generate(
    **inputs,
    max_new_tokens=50,  # 要生成的新token数量（约对应35-50个汉字）
    temperature=0.7,    # 控制随机性，值越低输出越确定
    top_p=0.9,
    do_sample=True,     # 开启采样生成，关闭则为贪心选概率最高的token
    repetition_penalty=1.1  # 降低重复内容的生成概率
)

# 解码并打印结果
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)

关键参数说明 #

• max_new_tokens：控制生成的文本长度，单位是token
• temperature：调整输出随机性，0表示完全确定，1表示随机性最高
• do_sample：设为False时会贪心选择概率最高的token，输出更稳定但易重复
• repetition_penalty：大于1的值会抑制重复内容生成

内容的提问来源于stack exchange，提问作者Franck Dernoncourt