Transformers documentation
Gemma 4 统一助手
该模型于 2026-06-03 贡献至 Hugging Face Transformers。
Gemma 4 统一助手
概述
Gemma 4 Unified Assistant 是一个小型的纯文本模型，它使用多 Token 预测（MTP）方法及其关联的候选生成器，为 Gemma 4 Unified 模型启用推测解码。我们提供了 Gemma 4 12B 模型 IT 变体的预训练模型。
有关更多信息，请参阅 Gemma4 Assistant。在架构和概念上，它们与其基础模型共享相同的概念，并存在相同的差异：
- 整个模型使用 KV 共享。这项技术最初在 Gemma 3n 中引入，允许模型复用由助手所支持的目标模型填充的 KV 缓存，从而使助手可以完全跳过 pre-fille 阶段，并显著降低前向传播过程中的注意力计算量。
position_ids的值是恒定的。由于 KV 缓存是共享的，且助手没有更新缓存的方式，因此助手会从相同的位置 ID 预测所有 token。- 输入是嵌入和隐藏状态的拼接。为了适配静态 KV 缓存和
position_ids，模型将输入视为目标模型最近一次看到的 token 的embedding与hidden_states的拼接，并通过nn.Linear变换将其投影到辅助模型的空间中。在辅助解码循环中，“最近一次看到的 token”的定义会不断变化。对于预填充后起草的第一个 token，最近一次看到的 token 是提示词中的最后一个 token。对于后续的起草步骤，最近一次看到的 token 是辅助模型生成的最后一个 token（在同一轮起草内），或是目标模型接受的最后一个 token（在不同轮起草之间）。 - 交叉注意力用于充分利用目标模型的上下文。交叉注意力允许助手生成的查询状态关注来自目标模型的共享 KV 缓存值，从而使助手在每一轮草稿生成中能够更准确地预测更多的草稿 token。
使用示例
以下示例演示了如何使用 Pipeline 或 AutoModel 类基于图像生成文本。
<hfoptions id="usage"> <hfoption id="Pipeline">import torch
from transformers import pipeline
pipeline = pipeline(
task="image-text-to-text",
model="google/gemma-4-12B-it",
assistant_model="google/gemma-4-12B-it-assistant",
)
pipeline(
images="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg",
text="<|image|>\n\nWhat is shown in this image?"
)import torch
from transformers import AutoProcessor, AutoModelForImageTextToText
model = AutoModelForImageTextToText.from_pretrained(
"google/gemma-4-12B-it",
dtype=torch.bfloat16,
device_map="auto",
)
assistant_model = AutoModelForCausalLM.from_pretrained(
"google/gemma-4-12B-it-assistant",
dtype=torch.bfloat16,
device_map="auto",
)
processor = AutoProcessor.from_pretrained(
"google/gemma-4-12B-it",
padding_side="left"
)
messages = [
{
"role": "user", "content": [
{"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"},
{"type": "text", "text": "What is shown in this image?"},
]
},
]
inputs = processor.apply_chat_template(
messages,
tokenize=True,
return_dict=True,
return_tensors="pt",
add_generation_prompt=True,
).to(model.device)
input_len = inputs["input_ids"].shape[-1]
output = model.generate(**inputs, max_new_tokens=50, assistant_model=assistant_model)
print(processor.decode(output[0][input_len:], skip_special_tokens=True))Gemma4UnifiedAssistantConfig
class transformers.Gemma4UnifiedAssistantConfig
< source >( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonetext_config: transformers.configuration_utils.PreTrainedConfig | dict[str, typing.Any] | None = Nonebackbone_hidden_size: int = 3840use_ordered_embeddings: bool = Falsenum_centroids: int = 2048centroid_intermediate_top_k: int = 32tie_word_embeddings: bool = True )
参数
- text_config (
Union[~configuration_utils.PreTrainedConfig, dict[str, Any]], 可选) — 文本骨干网络的配置对象或字典。 - backbone_hidden_size (
int，默认为 3840) — 该助手训练时所使用的目标统一模型的隐藏层大小。 - use_ordered_embeddings (
bool，默认为False) — 如果为True，则使用需要重新排序的最优顺序嵌入表。 - num_centroids (
int，默认为 2048) — 基于质心的掩码嵌入器中质心的总数。 - centroid_intermediate_top_k (
int，默认为 32) — 推理期间选中的活动质心数量。 - tie_word_embeddings (
bool，可选，默认为True) — 是否根据模型的tied_weights_keys映射来绑定权重嵌入。
这是用于存储 Gemma4 Unified AssistantModel 配置的类。它根据指定的参数实例化 Gemma4 Unified Assistant 模型，并定义模型架构。使用默认值实例化配置将生成与 google/gemma-4-12b-it 类似的配置。
配置对象继承自 PreTrainedConfig，可用于控制模型输出。有关更多信息，请参阅 PreTrainedConfig 的文档。
示例：
>>> from transformers import (
... Gemma4UnifiedAssistantConfig,
... Gemma4UnifiedAssistantForCausalLM,
... Gemma4UnifiedTextConfig,
... )
>>> # Initializing a Gemma 4 Unified Text config for the assistant.
>>> text_config = Gemma4UnifiedTextConfig(...)
>>> # Initializing a Gemma 4 Unified Assistant config.
>>> configuration = Gemma4UnifiedAssistantConfig(text_config=text_config)
>>> # Initializing a model from the configuration.
>>> model = Gemma4UnifiedAssistantForCausalLM(configuration)
>>> # Accessing the model configuration
>>> configuration = model.configGemma4UnifiedAssistantForCausalLM
class transformers.Gemma4UnifiedAssistantForCausalLM
< source >( config: Gemma4UnifiedAssistantConfig )
参数
- config (Gemma4UnifiedAssistantConfig) — 模型配置类，包含模型的所有参数。使用配置文件进行初始化时，不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。
用于 Gemma 4 的多 token 预测辅助解码模型。
该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。
该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。
为辅助模型准备注意力掩码；在此实例中，shared_kv_states 充当过去的缓存。
我们使用双向掩码来考虑因果性
- 对于
q_len == 1这种边界情况，无论采用哪种方式，其行为都等同于完整注意力机制，因此没有区别。 - 当
q_idx=1且kv>=1- 通过翻转 kv 轴，我们从面向未来的视角切换到面向过去的视角
- 为了处理与位置无关的填充，我们在初始创建之前还会翻转基础注意力掩码
forward
< source >( input_ids: typing.Optional[torch.Tensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: dict[str, torch.Tensor] | None = Noneshared_kv_states: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = Noneuse_cache: bool | None = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> )
参数
- input_ids (
torch.Tensor形状为(batch_size, sequence_length), 可选) — 输入序列中各 token 在词表中的索引。默认情况下，填充（padding）将被忽略。可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- inputs_embeds（形状为
(batch_size, sequence_length, hidden_size)的torch.Tensor，可选） — 可选地，您可以选择直接传入嵌入表示，而不是传入input_ids。如果您希望对input_ids索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权，这将非常有用。 - position_ids (
torch.LongTensor形状为(batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为[0, config.n_positions - 1]. - attention_mask (
dict[str, torch.Tensor]形状为(batch_size, sequence_length), 可选) — 用于避免对填充 token 索引执行注意力的掩码。掩码值取自[0, 1]:- 1 表示未被屏蔽的 token，
- 对于被掩码的 token，值为 0。
- shared_kv_states（
dict[str, torch.Tensor，形状为(batch_size, 1, q_len, kv_len)，可选） — 一个字典，包含此模型中每个layer_type最后一层计算得到的 KV 值。 - use_cache (
bool，可选) — 如果设置为True，则会返回past_key_values键值状态，可用于加速解码（参见past_key_values）。
Gemma4UnifiedAssistantForCausalLM 的 forward 方法重写了 __call__ 特殊方法。
尽管前向传播的实现需要在此函数中定义，但之后应调用
Module实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。
示例：
>>> from transformers import AutoTokenizer, Gemma4UnifiedAssistantForCausalLM, Gemma4UnifiedForConditionalGeneration
>>> model = Gemma4UnifiedForConditionalGeneration.from_pretrained("google/gemma-4-e2b-it")
>>> assistant_model = Gemma4UnifiedAssistantForCausalLM.from_pretrained("google/gemma-4-e2b-it-assistant")
>>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-4-e2b-it")
>>> prompt = "What is your favorite condiment?"
>>> inputs = tokenizer(prompt, return_tensors="pt")
>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, assistant_model=assistant_model, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True)[0]
"What is your favorite condiment?"