Transformers

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

该模型于 2026-05-05 贡献至 Hugging Face Transformers。

Gemma 4 助手

概述

Gemma 4 Assistant 是一个小型的纯文本模型，它使用多 Token 预测（MTP）方法及其关联的候选生成器，为 Gemma 4 模型启用推测解码。我们提供了 Gemma 4 E2B、E4B、31B 和 26B-A4B（MoE）模型 IT 变体的预训练模型。

在架构上，Gemma 4 Assistant 与其他 Gemma 4 模型共享相同的 Gemma4TextModel 主干网络，但在几个关键方面有所不同：

整个模型使用 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。

您可以在 Gemma 4 发布中找到所有原始的 Gemma 4 Assistant 检查点。

使用示例

以下示例演示了如何使用 Pipeline 或 AutoModel 类基于图像生成文本。

import torch
from transformers import pipeline

pipeline = pipeline(
    task="image-text-to-text",
    model="google/gemma-4-E2B-it",
    assistant_model="google/gemma-4-E2B-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?"
)

</hfoption> <hfoption id="AutoModel">

import torch
from transformers import AutoProcessor, AutoModelForImageTextToText

model = AutoModelForImageTextToText.from_pretrained(
    "google/gemma-4-E2B-it",
    dtype=torch.bfloat16,
    device_map="auto",
)
assistant_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-E2B-it-assistant",
    dtype=torch.bfloat16,
    device_map="auto",
)

processor = AutoProcessor.from_pretrained(
    "google/gemma-4-E2B-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))

</hfoption>

Gemma4AssistantConfig

class transformers.Gemma4AssistantConfig

< 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 = 1536use_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，默认为 1536) — 该助手训练时所用目标模型的隐藏层大小。
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 AssistantModel 配置的类。它根据指定的参数实例化 Gemma4 Assistant 模型，并定义模型架构。使用默认值实例化配置，将得到与 google/gemma-4-e2b-it 类似的配置。

配置对象继承自 PreTrainedConfig，可用于控制模型输出。有关更多信息，请参阅 PreTrainedConfig 的文档。

示例：

>>> from transformers import (
>>>     Gemma4AssistantConfig,
>>>     Gemma4AssistantForCausalLM,
>>>     Gemma4TextConfig,
>>> )

>>> # Initializing a Gemma 4 Text config similar to google/gemma-4-e2b-it-assistant.
>>> text_config = Gemma4TextConfig(...)

>>> # Initializing a Gemma 4 Assistant config similar to google/gemma-4-e2b-it-assistant.
>>> configuration = Gemma4AssistantConfig(text_config)

>>> # Initializing a model from the google/gemma-4-e2b-it-assistant configuration.
>>> model = Gemma4AssistantForCausalLM(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config

Gemma4AssistantForCausalLM

类 transformers.Gemma4AssistantForCausalLM

< source >

( config: Gemma4AssistantConfig )

参数

config (Gemma4AssistantConfig) — 模型配置类，包含模型的所有参数。使用配置文件进行初始化时，不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

用于 Gemma 4 的多 token 预测辅助解码模型。

该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

该模型也是 PyTorch torch.nn.Module 的子类。请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。

create_attention_masks

< source >

( inputs_embedsattention_maskshared_kv_states )

为辅助模型准备注意力掩码；在此实例中，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()。

什么是输入 ID？
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].
什么是位置 ID？
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）。

Gemma4AssistantForCausalLM 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义，但之后应调用 Module 实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。

示例：

>>> from transformers import AutoTokenizer, Gemma4AssistantForCausalLM, Gemma4ForCausalLM

>>> model = Gemma4ForCausalLM.from_pretrained("google/gemma-4-e2b-it")
>>> assistant_model = Gemma4AssistantForCausalLM.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?"

Update on GitHub

←Florence2 Gemma 4 统一助手→