# LLaDA2

[LLaDA2](https://huggingface.co/collections/inclusionAI/llada21) — это семейство дискретных диффузионных языковых моделей, генерирующих текст посредством итеративного уточнения по блокам. Вместо авторегрессивной генерации токенов по одному LLaDA2 начинает с полностью замаскированной последовательности и постепенно снимает маски с токенов по мере уверенности на нескольких шагах уточнения.

## Использование

```py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from diffusers import BlockRefinementScheduler, LLaDA2Pipeline

model_id = "inclusionAI/LLaDA2.1-mini"
model = AutoModelForCausalLM.from_pretrained(
    model_id, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
scheduler = BlockRefinementScheduler()

pipe = LLaDA2Pipeline(model=model, scheduler=scheduler, tokenizer=tokenizer)
output = pipe(
    prompt="Write a short poem about the ocean.",
    gen_length=256,
    block_length=32,
    num_inference_steps=32,
    threshold=0.7,
    editing_threshold=0.5,
    max_post_steps=16,
    temperature=0.0,
)
print(output.texts[0])
```

## Обработчики событий

Обработчики (callbacks) выполняются после каждого шага уточнения. Передайте `callback_on_step_end_tensor_inputs`, чтобы выбрать, какие тензоры будут включены в `callback_kwargs`. В текущей реализации предоставляются `block_x` (окно последовательности, которое уточняется) и `transfer_index` (маска фиксации заполнения маски); верните из обработчика `{"block_x": ...}`, чтобы заменить окно.

```py
def on_step_end(pipe, step, timestep, callback_kwargs):
    block_x = callback_kwargs["block_x"]
    # Inspect or modify `block_x` here.
    return {"block_x": block_x}

out = pipe(
    prompt="Write a short poem.",
    callback_on_step_end=on_step_end,
    callback_on_step_end_tensor_inputs=["block_x"],
)
```

## Рекомендуемые параметры

Модели LLaDA2.1 поддерживают два режима:

| Режим | `threshold` | `editing_threshold` | `max_post_steps` |
|------|-------------|---------------------|------------------|
| Качество | 0.7 | 0.5 | 16 |
| Скорость | 0.5 | `None` | 16 |

Передайте `editing_threshold=None`, `0.0` или отрицательное значение, чтобы отключить редактирование после маски.

Для моделей LLaDA2.0 отключите редактирование, передав `editing_threshold=None` или `0.0`.

Для всех моделей: `block_length=32`, `temperature=0.0`, `num_inference_steps=32`.

## LLaDA2Pipeline[[diffusers.LLaDA2Pipeline]]

#### diffusers.LLaDA2Pipeline[[diffusers.LLaDA2Pipeline]]

```python
diffusers.LLaDA2Pipeline(model: Any, scheduler: BlockRefinementScheduler, tokenizer: Any | None = None)
```

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/llada2/pipeline_llada2.py#L59)

Пайплайн для генерации дискретного диффузионного текста в стиле LLaDA2 с помощью блочного итеративного уточнения.

Этот конвейер поддерживает шаблонную последовательность, заполненную токеном `mask_token_id`, и уточняет её блоками. На каждом шаге уточнения он выбирает кандидата токены для активного блока и фиксирует подмножество на основе уверенности.

Ожидается, что модель будет принимать attention mask и `position_ids` и возвращать logits формы `[batch, seq, vocab_size]`.

#### __call__[[diffusers.LLaDA2Pipeline.__call__]]

```python
__call__(prompt: str | list[str] | None = None, messages: list[dict[str, str]] | None = None, input_ids: torch.LongTensor | None = None, attention_mask: torch.LongTensor | None = None, use_chat_template: bool = True, add_generation_prompt: bool = True, gen_length: int = 2048, block_length: int | None = None, num_inference_steps: int = 32, temperature: float = 0.0, top_p: float | None = None, top_k: int | None = None, sampling_method: str = 'multinomial', threshold: float = 0.7, editing_threshold: float | None = 0.5, max_post_steps: int = 16, minimal_topk: int = 1, eos_early_stop: bool = True, eos_token_id: int | None = None, mask_token_id: int | None = None, generator: torch.Generator | None = None, output_type: str = 'text', return_dict: bool = True, callback_on_step_end: Callable[[int, int, dict], None] | PipelineCallback | MultiPipelineCallbacks | None = None, callback_on_step_end_tensor_inputs: list[str] | None = None)
```

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/llada2/pipeline_llada2.py#L242)

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

prompt (`str` или `List[str]`, *необязательный*) : Текст подсказки. Если `use_chat_template` равно `True` (по умолчанию) и доступен токенизатор с шаблоном чата, подсказка оборачивается в сообщение чата перед токенизацией.

messages (`List[Dict[str, str]]`, *необязательный*) : Сообщения чата для кодирования (например, `[{"role": "user", "content": "Hello"}]`). Имеет приоритет над `prompt`, если указано. Требует токенизатора с `apply_chat_template`.

input_ids (`torch.LongTensor`, *необязательный*) : Предварительно токенизированные ID входных данных. Имеет приоритет над `prompt` и `messages`.

attention_mask (`torch.LongTensor`, *необязательный*) : Маска для каждого токена (1 для валидных токенов промпта, 0 для заполнителей), соответствующая форме `input_ids`. Используется только при передаче `input_ids`. Если параметр опущен (при наличии `input_ids`), все позиции считаются валидными. При формировании входных данных из `prompt` / `messages` маска токенизатора передаётся автоматически.

use_chat_template (`bool`, по умолчанию `True`) : Нужно ли оборачивать запрос в шаблон чата.

add_generation_prompt (`bool`, по умолчанию `True`) : добавлять ли промпт генерации при использовании шаблонов чата.

gen_length (`int`) : Количество токенов для генерации.

block_length (`int`, *optional*) : Размер блока для уточнения. Если не указан, используется `block_length`, настроенный в планировщике.

num_inference_steps (`int`) : Количество шагов уточнения на блок.

temperature (`float`) : Температура выборки.

top_p (`float`, *необязательный*) : Порог обрезки для ядерной выборки.

top_k (`int`, *optional*) : Порог выборки Top-k.

sampling_method (`str`) : Метод выборки (`auto`, `greedy`, `multinomial`).

threshold (`float`) : Порог уверенности для фиксации токенов.

editing_threshold (`float`, *optional*) : Порог уверенности для редактирования уже зафиксированных (не маскированных) токенов. Если значение положительное, после разрешения всех маскированных токенов в блоке конвейер продолжает уточнение: если модель предсказывает другой токен с уверенностью выше этого порога, существующий токен заменяется. Установите `None`, `0.0` или отрицательное значение, чтобы отключить редактирование. По умолчанию `0.5`.

max_post_steps (`int`) : Максимальное количество дополнительных итераций уточнения после разрешения всех токенов маски в блоке. Используется только при включённом `editing_threshold`. По умолчанию `16`.

minimal_topk (`int`) : Минимальное количество токенов для фиксации на каждом шаге.

eos_early_stop (`bool`) : Нужно ли останавливаться после фиксации EOS в блоке.

eos_token_id (`int`, *optional*) : ID токена EOS, используемого для раннего завершения.

mask_token_id (`int`, *необязательный*) : ID маскирующего токена, используемого в шаблоне.

generator (`torch.Generator`, *необязательный*) : генератор случайных чисел для выборки.

output_type (`str`, по умолчанию `"text"`) : Формат вывода. `"text"` декодирует последовательности в строки (требуется токенизатор). `"seq"` возвращает только исходные последовательности ID токенов.

return_dict (`bool`, *необязательный*, по умолчанию `True`) : Возвращать ли [LLaDA2PipelineOutput](/docs/diffusers/main/ru/api/pipelines/llada2#diffusers.LLaDA2PipelineOutput) вместо кортежа.

callback_on_step_end (`Callable` или `PipelineCallback`, *необязательный*) : Callback, выполняемый после каждого шага уточнения с сигнатурой `callback_on_step_end(self, step: int, timestep: int, callback_kwargs: Dict)`.

callback_on_step_end_tensor_inputs (`List[str]`, *необязательный*) : Ключи тензоров, передаваемые в callback. Допустимые ключи: `block_x`, `transfer_index`, `editing_transfer_index`, `sampled_tokens`, `sampled_probs`, `active_block`.

**Возвращает:** [LLaDA2PipelineOutput](/docs/diffusers/main/ru/api/pipelines/llada2#diffusers.LLaDA2PipelineOutput) или `tuple`

Если `return_dict` равно `True`, возвращается [LLaDA2PipelineOutput](/docs/diffusers/main/ru/api/pipelines/llada2#diffusers.LLaDA2PipelineOutput), иначе возвращается `tuple`, где первый элемент — это сгенерированные идентификаторы токенов (`torch.LongTensor`), а второй — декодированные тексты (`list[str]`), либо `None`, если `output_type` равно `"seq"`.

Генерация текста с покомпонентным уточнением.

Примеры:
```python
>>> import torch
>>> from transformers import AutoModelForCausalLM, AutoTokenizer
>>> from diffusers import BlockRefinementScheduler, LLaDA2Pipeline

>>> model_id = "inclusionAI/LLaDA2.1-mini"
>>> model = AutoModelForCausalLM.from_pretrained(
...     model_id, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto"
... )
>>> tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
>>> scheduler = BlockRefinementScheduler()

>>> pipe = LLaDA2Pipeline(model=model, scheduler=scheduler, tokenizer=tokenizer)
>>> output = pipe(prompt="What is the meaning of life?", gen_length=256)
>>> print(output.texts[0])
```

## LLaDA2PipelineOutput[[diffusers.LLaDA2PipelineOutput]]

#### diffusers.LLaDA2PipelineOutput[[diffusers.LLaDA2PipelineOutput]]

```python
diffusers.LLaDA2PipelineOutput(sequences: torch.LongTensor, texts: list[str] | None = None)
```

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/llada2/pipeline_llada2.py#L54)