Inference with Hugging Face Transformers

[Translation Pending]\n\n# 第 03 期 | Hugging Face Transformers 推理实战

🎯 学习目标

掌握使用 Hugging Face transformers 库加载 Google Gemma 模型进行推理的基本流程。
理解并对比 AutoModelForCausalLM 与 pipeline 两种推理方式的异同及适用场景。
学会如何配置 temperature 和 top_p 等采样参数，以控制生成文本的创造性和多样性。
实践如何实现流式输出（Streaming Output），提升用户体验。

📖 核心概念讲解

Google Gemma 模型作为 Google DeepMind 推出的先进开源大语言模型，其与 Hugging Face 生态的深度集成，使得开发者能够便捷地利用 transformers 库进行模型加载、推理及微调。本期教程将聚焦于推理实战，深入探讨 transformers 库中两种主要的推理范式：高层抽象的 pipeline 和底层灵活的 AutoModelForCausalLM。

3.1 Hugging Face `transformers` 库概览

transformers 库是 Hugging Face 的核心项目，提供了一个统一的 API 来加载、训练和使用预训练模型，涵盖了自然语言处理、计算机视觉、语音等多个领域。它支持 PyTorch、TensorFlow 和 JAX 等多种深度学习框架，极大地简化了大型模型的应用。对于 Gemma 而言，transformers 库提供了开箱即用的支持，包括模型架构、预训练权重和对应的分词器。

3.2 `pipeline`：高层抽象的推理利器

pipeline 是 transformers 库提供的一种高级抽象，旨在简化常见任务（如文本生成、文本分类、问答等）的使用。它封装了模型的加载、分词、前向传播和后处理等所有步骤，用户只需提供输入数据即可获得结果，无需关心底层细节。

3.2.1 工作原理

当您使用 pipeline 进行文本生成时，其内部通常会执行以下步骤：

模型与分词器加载： 根据指定的任务和模型名称，自动加载预训练模型（如 AutoModelForCausalLM）和对应的分词器（如 AutoTokenizer）。
输入处理： 使用分词器将输入的文本转换为模型可理解的 token ID 序列。
模型推理： 将 token ID 序列送入模型进行前向传播，生成下一个 token 的概率分布。
解码与采样： 根据预设的采样策略（如贪婪搜索、束搜索、Top-K、Top-P 或温度采样）从概率分布中选择下一个 token。
输出处理： 将生成的 token ID 序列解码回人类可读的文本。

3.2.2 适用场景

快速原型开发： 无需深入了解模型细节，即可快速验证想法。
简单任务： 当您的需求与 pipeline 预设的任务类型高度匹配时。
教学演示： 简洁明了，易于理解。

3.3 `AutoModelForCausalLM` 与 `AutoTokenizer`：底层灵活的控制

AutoModelForCausalLM 和 AutoTokenizer 是 transformers 库中用于加载模型和分词器的核心类。AutoModelForCausalLM 专门用于自回归（因果）语言模型，即根据前面的 token 预测下一个 token。AutoTokenizer 则负责将原始文本转换为模型可处理的数字 ID 序列，并将数字 ID 序列转换回文本。

3.3.1 工作原理

使用 AutoModelForCausalLM 和 AutoTokenizer 进行推理时，您需要手动控制以下步骤：

分词器加载： 使用 AutoTokenizer.from_pretrained() 加载与模型匹配的分词器。
模型加载： 使用 AutoModelForCausalLM.from_pretrained() 加载模型权重。
输入编码： 使用分词器将输入文本编码为 token ID 序列和注意力掩码。
生成： 调用模型的 generate() 方法，传入编码后的输入和各种生成参数（如 max_new_tokens、do_sample、temperature、top_p 等）。
输出解码： 将模型生成的 token ID 序列解码回文本。

3.3.2 适用场景

高级定制： 需要对分词、模型输入、生成策略进行细粒度控制时。
性能优化： 可以手动将模型和输入移动到特定设备（如 GPU），进行更精细的内存管理。
流式输出： 实现 token-by-token 的实时输出，提升用户体验。
集成到复杂系统： 作为更大系统的一部分，需要与其他组件紧密协作时。

3.3.3 对比图示

+---------------------+      +---------------------+
|      用户输入文本     |      |      用户输入文本     |
+----------+----------+      +----------+----------+
           |                            |
           |                            |
           v                            v
+---------------------+      +---------------------+
|    Hugging Face     |      |    Hugging Face     |
|       pipeline      |      |     AutoTokenizer   |
| (text-generation)   |      |                     |
|  - 自动加载模型/分词器  |      |  - 文本编码/解码器     |
|  - 自动处理输入/输出  |      +----------+----------+
|  - 封装生成逻辑       |                 | token IDs
+----------+----------+                 | Attention Mask
           |                            v
           |                            +---------------------+
           v                            |    Hugging Face     |
+---------------------+               |  AutoModelForCausalLM |
|     生成文本结果     |               |   - 加载模型权重      |
+---------------------+               |   - 执行生成过程      |
|     (高层抽象)      |               |     (generate())    |
+---------------------+               +----------+----------+
                                                 | token IDs
                                                 v
                                        +---------------------+
                                        |    Hugging Face     |
                                        |     AutoTokenizer   |
                                        |    (decode())       |
                                        +----------+----------+
                                                 |
                                                 v
                                        +---------------------+
                                        |     生成文本结果     |
                                        +---------------------+
                                        |     (底层控制)      |
                                        +---------------------+

3.4 采样参数调优：控制生成文本的风格

在语言模型生成文本时，除了最简单的贪婪搜索（每次选择概率最高的 token）和束搜索（Beam Search），我们还可以利用采样策略来引入随机性，使生成内容更加多样和富有创意。

3.4.1 `temperature` (温度)

作用： temperature 参数用于控制生成文本的随机性。它通过调整模型输出的概率分布的“锐度”来实现。
原理： temperature 值越高，模型输出的概率分布越平坦，导致模型在选择下一个 token 时有更多“不那么确定”的选项，从而生成更具创造性和多样性的文本。反之，temperature 值越低（接近 0），概率分布越尖锐，模型倾向于选择概率最高的 token，生成结果更确定、更保守，甚至可能重复。当 temperature=0 时，等同于贪婪搜索。
取值范围： 通常在 0.0 到 2.0 之间。

3.4.2 `top_p` (核采样/Nucleus Sampling)

作用： top_p 是一种更智能的采样策略，它在每次生成时，只考虑累积概率达到 p 的最小 token 集合。
原理： 假设模型输出了一个 token 概率分布，top_p 会从概率最高的 token 开始累加，直到累积概率达到 p。然后，它只在这个子集中进行采样。这确保了采样的 token 始终是“合理”的，同时仍然保持一定的多样性。
取值范围： 通常在 0.0 到 1.0 之间。top_p=1.0 意味着考虑所有 token，top_p=0.0 结合 do_sample=True 可能会导致不可预测的行为，通常与 temperature=0 结合使用时等同于贪婪搜索。

3.4.3 `do_sample`

作用： 一个布尔值参数，必须设置为 True 才能启用 temperature 和 top_p 等采样策略。如果 do_sample=False，则默认使用贪婪搜索或束搜索（如果 num_beams > 1）。

3.5 流式输出 (Streaming Output)

对于大型语言模型，生成完整响应可能需要数秒甚至数十秒。在此期间，用户界面如果没有任何反馈，会造成糟糕的用户体验。流式输出（即 token-by-token 地实时显示生成内容）可以显著提升用户体验，让用户感觉模型响应更快、更具交互性。

在 transformers 库中，实现流式输出通常涉及自定义回调或利用 TextStreamer 等工具。核心思想是模型在生成每个 token 后，立即将其解码并输出，而不是等待所有 token 生成完毕。

💻 实战演示

在开始实战之前，请确保您的环境已正确配置。

1. 安装必要的库：

pip install transformers accelerate bitsandbytes torch
# 注意：PyTorch 安装可能需要根据您的CUDA版本进行选择，
# 详见 PyTorch 官网：https://pytorch.org/get-started/locally/
# 例如，CUDA 12.1: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

2. 登录 Hugging Face：

Gemma 模型需要您在 Hugging Face 上接受其使用条款。登录后，您才能下载模型。

huggingface-cli login

按照提示输入您的 Hugging Face Access Token。

3. 选择 Gemma 模型：

本教程将使用 google/gemma-2b-it (Instruction-Tuned) 模型。对于本地运行，这是一个相对较小的模型，但仍建议使用 GPU 并启用 4 位量化以节省显存。如果您有更强大的 GPU (例如 24GB VRAM)，可以尝试 google/gemma-7b-it。

场景一：使用 `pipeline` 进行基本推理

pipeline 是最简单快捷的上手方式。

import torch
from transformers import pipeline

# 检查 CUDA 是否可用
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# 加载 Gemma 模型 (2B Instruction-Tuned)
# load_in_4bit=True 启用 4 位量化，可以显著减少显存占用，适合消费级 GPU
# torch_dtype=torch.bfloat16 可以提升性能，但需要支持 bfloat16 的 GPU
generator = pipeline(
    "text-generation",
    model="google/gemma-2b-it",
    model_kwargs={"torch_dtype": torch.bfloat16, "load_in_4bit": True},
    device=device
)

# Gemma 模型的指令格式
# 详细格式请参考：https://huggingface.co/google/gemma-2b-it#inference
prompt = "写一个关于人工智能如何改变教育的短故事。"
messages = [
    {"role": "user", "content": prompt},
]

# pipeline 默认会自动处理 prompt 格式
# 如果需要明确指定，可以使用 tokenizer.apply_chat_template
# max_new_tokens 控制生成文本的最大长度
# do_sample=True 开启采样，temperature 和 top_p 生效
# temperature 和 top_p 可以调整生成文本的创造性
outputs = generator(
    messages,
    max_new_tokens=256,
    do_sample=True,
    temperature=0.7,
    top_p=0.9
)

print("\n--- pipeline 基本推理结果 ---")
print(outputs[0]["generated_text"])

预期输出（部分）：

Using device: cuda
Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████

注意：

实际输出可能因模型版本、推理参数、随机性等因素而异。
首次运行会下载模型，可能需要一段时间。
Gemma-2B模型需要大约 5-6GB 显存 (4bit 量化后)。
如果您遇到显存不足的问题，请尝试关闭所有占用 GPU 资源的其他程序，或使用更小的模型（如将来可能出现的 Gemma-1B）。

场景二：使用 `AutoModelForCausalLM` 进行高级推理与参数调优

这种方式提供了对模型加载、分词和生成过程的更精细控制。我们将在此场景中重点演示 temperature 和 top_p 的效果。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer

# 检查 CUDA 是否可用
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# 1. 加载分词器和模型
model_id = "google/gemma-2b-it"
tokenizer = AutoTokenizer.from_pretrained(model_id)

# 加载模型，启用 4 位量化并移动到指定设备
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16, # 建议使用 bfloat16 提升性能
    load_in_4bit=True,          # 启用 4 位量化
    device_map="auto"           # 自动将模型加载到可用的设备
)
# 注意：当使用 device_map="auto" 时，模型会自动放置到 GPU，无需手动 .to(device)

print("\n--- AutoModelForCausalLM 高级推理 ---")

# 2. Gemma 模型的指令格式
# 确保使用正确的 chat template
chat = [
    {"role": "user", "content": "帮我写一首关于秋天的五言绝句。"},
]
prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
print(f"Formatted prompt:\n{prompt}")

# 3. 编码输入
input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)

# 4. 生成文本并调整采样参数

print("\n--- 尝试不同的采样参数 ---")

# 示例 1: 较低的 temperature (更保守，重复性可能高)
print("\n--- 采样参数：temperature=0.3, top_p=0.9 ---")
outputs_low_temp = model.generate(
    **input_ids,
    max_new_tokens=128,
    do_sample=True,
    temperature=0.3,
    top_p=0.9,
    pad_token_id=tokenizer.eos_token_id # 确保设置 pad_token_id
)
generated_text_low_temp = tokenizer.decode(outputs_low_temp[0], skip_special_tokens=True)
print(generated_text_low_temp)

# 示例 2: 较高的 temperature (更具创造性，多样性高)
print("\n--- 采样参数：temperature=0.9, top_p=0.9 ---")
outputs_high_temp = model.generate(
    **input_ids,
    max_new_tokens=128,
    do_sample=True,
    temperature=0.9,
    top_p=0.9,
    pad_token_id=tokenizer.eos_token_id
)
generated_text_high_temp = tokenizer.decode(outputs_high_temp[0], skip_special_tokens=True)
print(generated_text_high_temp)

# 示例 3: 较低的 top_p (更集中于高概率词)
print("\n--- 采样参数：temperature=0.7, top_p=0.5 ---")
outputs_low_top_p = model.generate(
    **input_ids,
    max_new_tokens=128,
    do_sample=True,
    temperature=0.7,
    top_p=0.5,
    pad_token_id=tokenizer.eos_token_id
)
generated_text_low_top_p = tokenizer.decode(outputs_low_top_p[0], skip_special_tokens=True)
print(generated_text_low_top_p)

预期输出（部分，每次运行可能不同）：

Using device: cuda
Formatted prompt:
<bos><start_of_turn>user
帮我写一首关于秋天的五言绝句。
<end_of_turn>
<start_of_turn>model

--- 尝试不同的采样参数 ---

--- 采样参数：temperature=0.3, top_p=0.9 ---
<bos><start_of_turn>user
帮我写一首关于秋天的五言绝句。
<end_of_turn>
<start_of_turn>model
秋风吹落叶，
寒露浸霜林。
雁字南飞远，
思乡入梦深。

--- 采样参数：temperature=0.9, top_p=0.9 ---
<bos><start_of_turn>user
帮我写一首关于秋天的五言绝句。
<end_of_turn>
<start_of_turn>model
秋风起叶落，
霜染枫林红。
雁阵归南去，
独坐思故乡。

--- 采样参数：temperature=0.7, top_p=0.5 ---
<bos><start_of_turn>user
帮我写一首关于秋天的五言绝句。
<end_of_turn>
<start_of_turn>model
秋色入画来，
枫叶染霜红。
雁字排云去，
寒山独钓翁。

分析： 通过对比不同 temperature 和 top_p 参数下的生成结果，您可以观察到它们对文本创造性和多样性的影响。较高的 temperature 和 top_p 往往能生成更意想不到、更富诗意的表达，而较低的值则倾向于更常见、更保守的词语组合。

场景三：实现流式输出

流式输出是提升用户体验的关键。Hugging Face transformers 库提供了 TextStreamer 类，可以方便地实现这一功能。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer

# 检查 CUDA 是否可用
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# 1. 加载分词器和模型
model_id = "google/gemma-2b-it"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    load_in_4bit=True,
    device_map="auto"
)

print("\n--- 流式输出演示 ---")

# 2. Gemma 模型的指令格式
chat = [
    {"role": "user", "content": "请详细解释一下量子计算的基本原理和它可能带来的影响。"},
]
prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
print(f"Formatted prompt:\n{prompt}")

# 3. 编码输入
input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)

# 4. 初始化 TextStreamer
# TextStreamer 会在每个 token 生成后立即解码并打印
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

# 5. 生成文本，使用 streamer
print("\n--- 模型正在生成 (流式输出)... ---")
_ = model.generate(
    **input_ids,
    max_new_tokens=512,
    do_sample=True,
    temperature=0.7