# Gemma 4 Unified Assistant

## Обзор

Gemma 4 Unified Assistant — это небольшая текстовая модель, которая обеспечивает спекулятивное декодирование для моделей Gemma 4 Unified с использованием метода Multi-Token Prediction (MTP) и связанного с ним генератора кандидатов. Предоставляются предобученные модели для IT-вариантов модели Gemma 4 12B.

Для получения дополнительной информации см. [Gemma4 Assistant](./gemma4_assistant). Архитектурно и концептуально они разделяют одну и ту же концепцию и отличия от своей базовой модели:

*   **Вся модель использует общий KV-кэш**. Эта техника, впервые представленная вместе с [Gemma 3n](./gemma3n), позволяет модели повторно использовать KV-кэш, заполненный целевой моделью, которую поддерживает ассистент. Благодаря этому ассистент может полностью пропустить фазу pre-fille, что значительно снижает вычислительные затраты на внимание (attention) при прямом проходе.
*   **Значения `position_ids` остаются постоянными**. Поскольку кэш KV используется совместно, а ассистент не имеет способа обновить кэш, он предсказывает все токены из одного и того же идентификатора позиции.
*   **Входные данные представляют собой конкатенацию эмбеддингов и скрытых состояний**. Для адаптации к статическому KV-кэшу и `position_ids` модель принимает на вход конкатенацию `embedding` и `hidden_states` последнего просмотренного токена целевой модели и проецирует их в пространство модели-ассистента с помощью преобразования `nn.Linear`. Определение последнего просмотренного токена меняется в ходе цикла ассистированного декодирования. Для первого токена, сгенерированного после pre-fill, последним просмотренным токеном будет последний токен из промпта. Для последующих шагов генерации последним просмотренным токеном будет последний токен, сгенерированный ассистентом (в рамках раунда генерации) или последний токен, принятый целевой моделью (между раундами генерации).
*   **Кросс-внимание используется для максимального использования контекста целевой модели**. Кросс-внимание позволяет состояниям запросов, сгенерированным ассистентом, обращаться к общим значениям KV-кэша целевой модели, что позволяет ассистенту точно предсказывать больше токенов за каждый раунд черновика.

## Примеры использования

Приведённый ниже пример показывает, как генерировать текст на основе изображения с помощью [Pipeline](/docs/transformers/main/ru/main_classes/pipelines#transformers.Pipeline) или класса [AutoModel](/docs/transformers/main/ru/model_doc/auto#transformers.AutoModel).

```py
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?"
)
```

```py
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[[transformers.Gemma4UnifiedAssistantConfig]]

#### transformers.Gemma4UnifiedAssistantConfig[[transformers.Gemma4UnifiedAssistantConfig]]

```python
transformers.Gemma4UnifiedAssistantConfig(transformers_version: str | None = None, architectures: list[str] | None = None, output_hidden_states: bool | None = False, return_dict: bool | None = True, dtype: str | torch.dtype | None = None, chunk_size_feed_forward: int = 0, is_encoder_decoder: bool = False, id2label: dict[int, str] | dict[str, str] | None = None, label2id: dict[str, int] | dict[str, str] | None = None, problem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = None, text_config: transformers.configuration_utils.PreTrainedConfig | dict[str, typing.Any] | None = None, backbone_hidden_size: int = 3840, use_ordered_embeddings: bool = False, num_centroids: int = 2048, centroid_intermediate_top_k: int = 32, tie_word_embeddings: bool = True)
```

[Источник](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma4_unified_assistant/configuration_gemma4_unified_assistant.py#L25)

**Параметры:**

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](https://huggingface.co/google/gemma-4-12b-it)

Объекты конфигурации наследуются от [PreTrainedConfig](/docs/transformers/main/ru/main_classes/configuration#transformers.PreTrainedConfig) и могут использоваться для управления выводами модели. Подробнее см. в документации [PreTrainedConfig](/docs/transformers/main/ru/main_classes/configuration#transformers.PreTrainedConfig).

Пример:

```python
>>> 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.config
```

## Gemma4UnifiedAssistantForCausalLM[[transformers.Gemma4UnifiedAssistantForCausalLM]]

#### transformers.Gemma4UnifiedAssistantForCausalLM[[transformers.Gemma4UnifiedAssistantForCausalLM]]

```python
transformers.Gemma4UnifiedAssistantForCausalLM(config: Gemma4UnifiedAssistantConfig)
```

[Источник](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma4_unified_assistant/modeling_gemma4_unified_assistant.py#L115)

**Параметры:**

config ([Gemma4UnifiedAssistantConfig](/docs/transformers/main/ru/model_doc/gemma4_unified_assistant#transformers.Gemma4UnifiedAssistantConfig)) : Класс конфигурации модели со всеми параметрами модели. Инициализация с файлом конфигурации не загружает веса, связанные с моделью, только конфигурацию. Используйте метод [from_pretrained()](/docs/transformers/main/ru/main_classes/model#transformers.PreTrainedModel.from_pretrained) для загрузки весов модели.

Модель для вспомогательного декодирования на основе предсказания нескольких токенов с Gemma 4.

Эта модель наследуется от [PreTrainedModel](/docs/transformers/main/ru/main_classes/model#transformers.PreTrainedModel). Общие методы, которые библиотека реализует для всех своих моделей (например, загрузка или сохранение, изменение размера входных эмбеддингов, обрезка головок и т. д.), описаны в документации родительского класса.

Эта модель также является подклассом [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) в PyTorch. Используйте её как обычный модуль PyTorch и обращайтесь к документации PyTorch по всем вопросам, связанным с общим использованием и поведением.

#### create_attention_masks[[transformers.Gemma4UnifiedAssistantForCausalLM.create_attention_masks]]

```python
create_attention_masks(inputs_embeds, attention_mask, shared_kv_states)
```

[Источник](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma4_unified_assistant/modeling_gemma4_unified_assistant.py#L211)

Подготовьте маски внимания для модели-ассистента; в данном случае `shared_kv_states` выполняет роль кэша прошлых значений.

Мы используем двунаправленные маски для учёта причинно-следственных связей
- Для граничного случая `q_len == 1` разницы нет, так как он работает как полное внимание в любом случае.
- SWA интерпретирует окно как направленное в будущее, когда `q_idx=1` и `kv>=1`
  - Мы меняем перспективу с будущей на прошедшую, инвертируя ось kv
  - Чтобы учесть позиционно-инвариантную отсечку, мы также инвертируем базовую маску внимания перед её первоначальным созданием

#### forward[[transformers.Gemma4UnifiedAssistantForCausalLM.forward]]

```python
forward(input_ids: typing.Optional[torch.Tensor] = None, inputs_embeds: typing.Optional[torch.Tensor] = None, position_ids: typing.Optional[torch.LongTensor] = None, attention_mask: dict[str, torch.Tensor] | None = None, shared_kv_states: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None, use_cache: bool | None = None, **kwargs: *<class 'transformers.utils.generic.TransformersKwargs'>)
```

[Источник](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma4_unified_assistant/modeling_gemma4_unified_assistant.py#L140)

**Параметры:**

input_ids (`torch.Tensor` формы `(batch_size, sequence_length)`, *необязательный*) : Индексы токенов входной последовательности в словаре. По умолчанию заполнение (padding) игнорируется. Индексы можно получить с помощью [AutoTokenizer](/docs/transformers/main/ru/model_doc/auto#transformers.AutoTokenizer). Подробности см. в [PreTrainedTokenizer.encode()](/docs/transformers/main/ru/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.encode) и [PreTrainedTokenizer.__call__()](/docs/transformers/main/ru/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.__call__). [Что такое входные идентификаторы?](../glossary#input-ids)

inputs_embeds (`torch.Tensor` формы `(batch_size, sequence_length, hidden_size)`, *необязательно*) : По желанию, вместо передачи `input_ids` можно напрямую передать встроенное представление. Это полезно, если требуется больший контроль над тем, как индексы `input_ids` преобразуются в соответствующие векторы, чем это обеспечивает внутренняя матрица поиска встраиваний модели.

position_ids (`torch.LongTensor` формы `(batch_size, sequence_length)`, *необязательно*) : Индексы позиций токенов каждого входного последовательности в позиционных эмбеддингах. Выбираются из диапазона `[0, config.n_positions - 1]`.  [Что такое position IDs?](../glossary#position-ids)

attention_mask (`dict[str, torch.Tensor]` формы `(batch_size, sequence_length)`, *необязательный*) : Маска для предотвращения применения внимания к индексам токенов отсечки. Значения маски выбираются из `[0, 1]`:  - 1 для токенов, которые **не замаскированы**, - 0 для токенов, которые **замаскированы**.  [Что такое маски внимания?](../glossary#attention-mask)

shared_kv_states (`dict[str, torch.Tensor` формы `(batch_size, 1, q_len, kv_len)`, *необязательный*) : Словарь, содержащий вычисленные значения KV для последнего слоя каждого `layer_type` в этой модели.

use_cache (`bool`, *необязательный*) : Если установлено в `True`, состояния ключей и значений `past_key_values` возвращаются и могут использоваться для ускорения декодирования (см. `past_key_values`).

**Возвращаемое значение:**

`tuple[torch.Tensor, torch.Tensor]`

Метод forward класса [Gemma4UnifiedAssistantForCausalLM](/docs/transformers/main/ru/model_doc/gemma4_unified_assistant#transformers.Gemma4UnifiedAssistantForCausalLM) переопределяет специальный метод `__call__`.

Хотя рецепт для прямого прохода должен быть определён внутри этой функции, после этого следует вызывать экземпляр `Module`, а не эту функцию, поскольку первый выполняет шаги предварительной и постобработки, тогда как второй молча их игнорирует.

Пример:

```python
>>> 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?"
```