# LLaDA2

[LLaDA2](https://huggingface.co/collections/inclusionAI/llada21) 是一组离散扩散语言模型，通过分块迭代优化来生成文本。与自回归的逐 token 生成不同，LLaDA2 从一个完全被掩码的序列开始，并在多个优化步骤中根据置信度逐步解除 token 的掩码。

## 用法

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

## 回调

回调函数会在每个细化步骤之后运行。通过传入 `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/v0.43.0/src/diffusers/pipelines/llada2/pipeline_llada2.py#L59)

用于 LLaDA2 风格离散扩散文本生成的流水线，通过分块迭代优化实现。

该流水线维护一个填充了 `mask_token_id` 的模板序列，并分块对其进行细化。在每个细化步骤中，它会为当前活动块采样候选 token，并根据置信度提交其中的一部分。

该模型预期会接受注意力掩码和 `position_ids`，并返回形状为 `[batch, seq, vocab_size]` 的 logits。

#### __调用__[[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/v0.43.0/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` 的 tokenizer。

input_ids (`torch.LongTensor`, *可选*) : 预分词化的输入 ID。其优先级高于 `prompt` 和 `messages`。

attention_mask (`torch.LongTensor`, *可选*) : 逐 token 掩码（有效提示 token 为 1，填充为 0），其形状与 `input_ids` 一致。仅在提供 `input_ids` 时使用。若省略（且已给出 `input_ids`），则所有位置均视为有效。从 `prompt` / `messages` 构建输入时，分词器的掩码会自动传递。

use_chat_template (`bool`，默认为 `True`)：是否将提示词包裹在聊天模板中。

add_generation_prompt (`bool`，默认为 `True`)：在使用聊天模板时，是否添加生成提示。

gen_length (`int`) : 要生成的 token 数量。

block_length (`int`, *可选*) : 细化操作的块大小。如果未提供，则使用调度器配置的 `block_length`。

num_inference_steps (`int`) : 每个块的细化步数。

temperature (`float`) : 采样温度。

top_p (`float`, *可选*)：核采样截断值。

top_k (`int`，*可选*)：Top-k 采样截断值。

sampling_method (`str`) : 采样方法（`auto`、`greedy`、`multinomial`）。

threshold (`float`) : 提交 token 的置信度阈值。

editing_threshold (`float`, *可选*) : 用于编辑已提交（非掩码）token 的置信度阈值。当该值为正数时，在块中所有掩码 token 解析完成后，流水线将继续进行细化：如果模型预测的另一个 token 的置信度高于此阈值，则会替换现有的 token。将其设置为 `None`、`0.0` 或负值以禁用编辑功能。默认值为 `0.5`。

max_post_steps (`int`) : 当块中所有掩码 token 均被解析后，允许执行的最大额外细化迭代次数。仅在启用 `editing_threshold` 时使用。默认值为 `16`。

minimal_topk (`int`) : 每个步骤提交的最小 token 数量。

eos_early_stop (`bool`) : 是否在一个块中提交 EOS 后停止。

eos_token_id (`int`, *可选*) : 用于提前停止的 EOS token ID。

mask_token_id (`int`, *可选*) : 用于模板的掩码 token ID。

generator (`torch.Generator`，*可选*)：用于采样的随机数生成器。

output_type (`str`，默认为 `"text"`)：输出格式。`"text"` 将序列解码为字符串（需要分词器）。`"seq"` 仅返回原始 token ID 序列。

return_dict (`bool`，*可选*，默认为 `True`)：是否返回 [LLaDA2PipelineOutput](/docs/diffusers/v0.43.0/zh/api/pipelines/llada2#diffusers.LLaDA2PipelineOutput) 而不是元组。

callback_on_step_end (`Callable` 或 `PipelineCallback`，*可选*)：在每次细化步骤后执行的回调，其签名为 `callback_on_step_end(self, step: int, timestep: int, callback_kwargs: Dict)`。

callback_on_step_end_tensor_inputs (`List[str]`, *可选*) : 传递给回调函数的 Tensor 键。允许的键： `block_x`、`transfer_index`、`editing_transfer_index`、`sampled_tokens`、`sampled_probs`、`active_block`。

**返回：** [LLaDA2PipelineOutput](/docs/diffusers/v0.43.0/zh/api/pipelines/llada2#diffusers.LLaDA2PipelineOutput) 或 `tuple`

如果 `return_dict` 为 `True`，则返回 [LLaDA2PipelineOutput](/docs/diffusers/v0.43.0/zh/api/pipelines/llada2#diffusers.LLaDA2PipelineOutput)；否则返回一个 `tuple`，其中第一个元素是生成的 token ID（`torch.LongTensor`），第二个元素是解码后的文本（`list[str]`）；当 `output_type` 为 `"seq"` 时，第二个元素为 `None`。

使用分块细化生成文本。

示例：
```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/v0.43.0/src/diffusers/pipelines/llada2/pipeline_llada2.py#L54)