Transformers documentation

Gemma 4 Assistant

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Este modelo foi contribuído para o Hugging Face Transformers em 2026-05-05.

Gemma 4 Assistant

Visão geral

O Gemma 4 Assistant é um modelo pequeno, somente de texto, que habilita a decodificação especulativa para modelos Gemma 4 usando o método de Previsão de Múltiplos Tokens (MTP) e o gerador de candidatos associado. Modelos pré-treinados são fornecidos para as variantes IT dos modelos Gemma 4 E2B, E4B, 31B e 26B-A4B (MoE).

Arquiteturalmente, o Gemma 4 Assistant compartilha o mesmo backbone Gemma4TextModel que os outros modelos Gemma 4, mas difere em alguns pontos-chave:

  • O modelo inteiro usa compartilhamento de KV. Essa técnica, introduzida originalmente com o Gemma 3n, permite que o modelo reutilize o cache de KV preenchido pelo modelo de destino que o assistente suporta, permitindo que o assistente pule completamente a fase de pre-fille e reduza consideravelmente o cálculo de atenção durante a passagem direta.
  • Os valores de position_ids são constantes. Como o cache KV é compartilhado e o assistente não possui uma forma de atualizar o cache, ele prevê todos os tokens a partir do mesmo ID de posição.
  • As entradas são a concatenação dos embeddings e dos hidden states. Para se adaptar ao KV cache estático e a position_ids, o modelo recebe suas entradas como a concatenação do embedding e dos hidden_states do último token visto pelo modelo de destino e os projeta para o espaço do modelo assistente com uma transformação nn.Linear. A definição de último token visto muda ao longo do loop de decodificação assistida. Para o primeiro token rascunhado após o pre-fill, o último token visto será o último token do prompt. Para as etapas de rascunho subsequentes, o último token visto será o último token gerado pelo assistente (dentro de uma rodada de rascunho) ou o último token aceito pelo modelo de destino (entre as rodadas de rascunho).
  • A cross-attention é usada para aproveitar ao máximo o contexto do modelo de destino. A cross-attention permite que os estados de consulta gerados pelo assistente prestem atenção aos valores do cache KV compartilhados pelo modelo de destino, permitindo que o assistente preveja com precisão mais tokens rascunhados por rodada de rascunho.

Você pode encontrar todos os checkpoints originais do Gemma 4 Assistant na liberação Gemma 4.

Exemplos de uso

O exemplo abaixo demonstra como gerar texto com base em uma imagem usando Pipeline ou a classe AutoModel.

<hfoptions id="usage"> <hfoption id="Pipeline">
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

classe transformers.Gemma4AssistantConfig

< >

( 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 )

Parâmetros

  • text_config (Union[~configuration_utils.PreTrainedConfig, dict[str, Any]], opcional) — O objeto de configuração ou dicionário do backbone de texto.
  • backbone_hidden_size (int, padrão 1536) — Tamanho oculto do modelo de destino com o qual este assistente foi treinado.
  • use_ordered_embeddings (bool, padrão: False) — Se True, usa uma tabela de embeddings ordenada para o desempenho ótimo do modelo assistente, que precisa ser reordenada para alinhá-la ao modelo principal.
  • num_centroids (int, padrão 2048) — O número total de centróides.
  • centroid_intermediate_top_k (int, padrão 32) — O número de centroides ativos.
  • tie_word_embeddings (bool, opcional, padrão True) — Se deve vincular as embeddings de pesos de acordo com o mapeamento tied_weights_keys do modelo.

Esta é a classe de configuração para armazenar a configuração de um Gemma4 AssistantModel. Ela é usada para instanciar um modelo Gemma4 Assistant de acordo com os argumentos especificados, definindo a arquitetura do modelo. Instanciar uma configuração com os padrões resultará em uma configuração semelhante à do google/gemma-4-e2b-it

Os objetos de configuração herdam de PreTrainedConfig e podem ser usados para controlar as saídas do modelo. Consulte a documentação de PreTrainedConfig para mais informações.

Exemplo:

>>> 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

classe transformers.Gemma4AssistantForCausalLM

< >

( config: Gemma4AssistantConfig )

Parâmetros

  • config (Gemma4AssistantConfig) — Classe de configuração do modelo com todos os parâmetros do modelo. A inicialização com um arquivo de configuração não carrega os pesos associados ao modelo, apenas a configuração. Consulte o método from_pretrained() para carregar os pesos do modelo.

Um modelo para decodificação assistida baseada em previsão de múltiplos tokens com Gemma 4.

Este modelo herda de PreTrainedModel. Consulte a documentação da classe pai para os métodos genéricos que a biblioteca implementa para todos os seus modelos (como download ou salvamento, redimensionamento dos embeddings de entrada, poda de cabeças, etc.)

Este modelo também é uma subclasse de torch.nn.Module do PyTorch. Use-o como um Módulo PyTorch regular e consulte a documentação do PyTorch para todos os assuntos relacionados ao uso geral e ao comportamento.

create_attention_masks

< >

( inputs_embedsattention_maskshared_kv_states )

Prepare as máscaras de atenção para o modelo assistido; o shared_kv_states atua como cache passado nesta instância.

Usamos máscaras bidirecionais para levar em conta a causalidade

  • Não há diferença para o caso extremo de q_len == 1, pois ele atua como atenção completa, independentemente do contexto.
  • A SWA interpreta a janela como voltada para o futuro quando q_idx=1 e. Alternamos de uma perspectiva futura para uma perspectiva passada invertendo no eixo kv. Para levar em conta o padding invariante em relação à posição, também invertemos a máscara de atenção base antes da criação inicial. kv>=1
    • Alternamos de uma perspectiva futura para uma perspectiva passada ao inverter o eixo kv
    • Para levar em conta o padding invariante em relação à posição, também invertemos a máscara de atenção base antes da criação inicial

forward

< >

( 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'> )

Parâmetros

  • input_ids (torch.Tensor de formato (batch_size, sequence_length), opcional) — Índices dos tokens da sequência de entrada no vocabulário. O padding será ignorado por padrão.

    Os índices podem ser obtidos usando AutoTokenizer. Consulte PreTrainedTokenizer.encode() e PreTrainedTokenizer.call() para mais detalhes.

    O que são IDs de entrada?

  • inputs_embeds (torch.Tensor de forma (batch_size, sequence_length, hidden_size), opcional) — Opcionalmente, em vez de passar input_ids, você pode escolher passar diretamente uma representação incorporada. Isso é útil se você quiser ter mais controle sobre como converter os índices de input_ids nos vetores associados do que a matriz de busca de incorporação interna do modelo.
  • position_ids (torch.LongTensor de forma (batch_size, sequence_length), opcional) — Índices das posições dos tokens de cada sequência de entrada nos embeddings de posição. Selecionados no intervalo [0, config.n_positions - 1].

    O que são IDs de posição?

  • attention_mask (dict[str, torch.Tensor] de forma (batch_size, sequence_length), opcional) — Máscara para evitar a execução de atenção nos índices de tokens de preenchimento. Valores de máscara selecionados em [0, 1]:
    • 1 para tokens que não estão mascarados,
    • 0 para tokens que estão mascarados.

    O que são máscaras de atenção?

  • shared_kv_states (dict[str, torch.Tensor de forma (batch_size, 1, q_len, kv_len), opcional) — Um dicionário contendo os valores de KV calculados para a última camada de cada layer_type neste modelo.
  • use_cache (bool, opcional) — Se definido como True, os estados de chave-valor past_key_values são retornados e podem ser usados para acelerar a decodificação (veja past_key_values).

O método forward de Gemma4AssistantForCausalLM substitui o método especial __call__.

Embora a receita para a passagem direta precise ser definida dentro desta função, deve-se chamar a instância de Module em seguida, em vez desta, pois a primeira cuida da execução das etapas de pré e pós-processamento, enquanto a segunda as ignora silenciosamente.

Exemplo:

>>> 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