# LLaDA2

[LLaDA2](https://huggingface.co/collections/inclusionAI/llada21) é uma família de modelos de linguagem de difusão discreta que gera texto por meio de refinamento iterativo por blocos. Em vez da geração autoregressiva token a token, o LLaDA2 começa com uma sequência totalmente mascarada e desmascara progressivamente os tokens por confiança ao longo de múltiplas etapas de refinamento.

## Uso

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

Os callbacks são executados após cada etapa de refinamento. Passe `callback_on_step_end_tensor_inputs` para selecionar quais tensores são incluídos em `callback_kwargs`. Na implementação atual, `block_x` (a janela de sequência sendo refinada) e `transfer_index` (máscara de commit para preenchimento de máscaras) são fornecidos; retorne `{"block_x": ...}` do callback para substituir a janela.

```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"],
)
```

## Parâmetros recomendados

Os modelos LLaDA2.1 suportam dois modos:

| Modo | `threshold` | `editing_threshold` | `max_post_steps` |
|------|-------------|---------------------|------------------|
| Qualidade | 0.7 | 0.5 | 16 |
| Velocidade | 0.5 | `None` | 16 |

Pass `editing_threshold=None`, `0.0` ou um valor negativo para desativar a edição pós-máscara.

Para os modelos LLaDA2.0, desative a edição passando `editing_threshold=None` ou `0.0`.

Para todos os modelos: `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)
```

[Fonte](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/llada2/pipeline_llada2.py#L59)

Pipeline para geração de texto por difusão discreta no estilo LLaDA2, por meio de refinamento iterativo por blocos.

Este pipeline mantém uma sequência de modelo preenchida com `mask_token_id` e a refina em blocos. Em cada etapa de refinamento, ele amostra tokens candidatos para o bloco ativo e confirma um subconjunto com base na confiança.

O modelo deve aceitar uma máscara de atenção e `position_ids`, e retornar logits com a forma `[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)
```

[Fonte](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/llada2/pipeline_llada2.py#L242)

**Parâmetros:**

prompt (`str` ou `List[str]`, *opcional*) : Texto do prompt. Quando `use_chat_template` é `True` (padrão) e um tokenizer com um modelo de chat está disponível, o prompt é envolvido em uma mensagem de chat antes da tokenização.

messages (`List[Dict[str, str]]`, *opcional*) : Mensagens de chat para codificar (por exemplo, `[{"role": "user", "content": "Hello"}]`). Tem precedência sobre `prompt` quando fornecido. Requer um tokenizer com `apply_chat_template`.

input_ids (`torch.LongTensor`, *opcional*) : IDs de entrada pré-tokenizados. Tem precedência sobre `prompt` e `messages`.

attention_mask (`torch.LongTensor`, *opcional*) : Máscara por token (1 para tokens de prompt válidos, 0 para padding) que corresponde à forma de `input_ids`. Usado apenas quando `input_ids` é fornecido. Quando omitido (e `input_ids` é fornecido), todas as posições são tratadas como válidas. Ao construir entradas a partir de `prompt` / `messages`, a máscara do tokenizer é transmitida automaticamente.

use_chat_template (`bool`, padrão `True`) : Se deve envolver o prompt em um modelo de chat.

add_generation_prompt (`bool`, padrão `True`) : Se deve adicionar o prompt de geração ao usar modelos de chat.

gen_length (`int`) : Número de tokens a serem gerados.

block_length (`int`, *opcional*) : Tamanho do bloco para refinamento. Se não for fornecido, o `block_length` configurado do agendador será usado.

num_inference_steps (`int`) : Número de etapas de refinamento por bloco.

temperature (`float`) : Temperatura de amostragem.

top_p (`float`, *opcional*) : Corte de amostragem por núcleo.

top_k (`int`, *opcional*) : Corte de amostragem Top-k.

sampling_method (`str`) : Método de amostragem (`auto`, `greedy`, `multinomial`).

threshold (`float`) : Limiar de confiança para confirmar tokens.

editing_threshold (`float`, *opcional*) : Limiar de confiança para edição de tokens já comprometidos (não-máscara). Quando positivo, após todos os tokens de máscara em um bloco serem resolvidos, o pipeline continua refinando: se o modelo prever um token diferente com confiança acima deste limiar, o token existente é substituído. Defina como `None`, `0.0` ou um valor negativo para desabilitar a edição. O padrão é `0.5`.

max_post_steps (`int`) : Número máximo de iterações adicionais de refinamento após todos os tokens de máscara em um bloco serem resolvidos. Usado apenas quando `editing_threshold` está habilitado. O valor padrão é `16`.

minimal_topk (`int`) : Número mínimo de tokens a serem confirmados por etapa.

eos_early_stop (`bool`) : Se deve parar após o commit do EOS em um bloco.

eos_token_id (`int`, *opcional*) : ID do token EOS a ser usado para parada antecipada.

mask_token_id (`int`, *opcional*) : ID do token de máscara a ser usado no modelo.

generator (`torch.Generator`, *opcional*) : RNG para amostragem.

output_type (`str`, padrão `"text"`) : Formato de saída. `"text"` decodifica sequências em strings (requer um tokenizer). `"seq"` retorna apenas sequências brutas de IDs de tokens.

return_dict (`bool`, *opcional*, padrão `True`) : Se deve retornar um [LLaDA2PipelineOutput](/docs/diffusers/main/pt/api/pipelines/llada2#diffusers.LLaDA2PipelineOutput) em vez de uma tupla.

callback_on_step_end (`Callable` ou `PipelineCallback`, *opcional*) : Callback executado após cada etapa de refinamento com a assinatura `callback_on_step_end(self, step: int, timestep: int, callback_kwargs: Dict)`.

callback_on_step_end_tensor_inputs (`List[str]`, *opcional*) : Chaves de tensor para passar ao callback. Chaves permitidas: `block_x`, `transfer_index`, `editing_transfer_index`, `sampled_tokens`, `sampled_probs`, `active_block`.

**Retorna:** [LLaDA2PipelineOutput](/docs/diffusers/main/pt/api/pipelines/llada2#diffusers.LLaDA2PipelineOutput) ou `tuple`

Se `return_dict` for `True`, [LLaDA2PipelineOutput](/docs/diffusers/main/pt/api/pipelines/llada2#diffusers.LLaDA2PipelineOutput) é retornado; caso contrário, uma `tuple` é retornada, em que o primeiro elemento são os IDs dos tokens gerados (`torch.LongTensor`) e o segundo elemento são os textos decodificados (`list[str]`), ou `None` quando `output_type` é `"seq"`.

Gere texto com refinamento por bloco.

Exemplos:
```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)
```

[Fonte](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/llada2/pipeline_llada2.py#L54)