Hugging Face's logo

Diffusers documentation

LLaDA2

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

LLaDA2

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

用法

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": ...} 即可替换该窗口。

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 模型支持两种模式：

模式thresholdediting_thresholdmax_post_steps
质量0.70.516
速度0.5None16

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

< >

( model: Anyscheduler: BlockRefinementSchedulertokenizer: Any | None = None )

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

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

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

__call__

< >

( prompt: str | list[str] | None = Nonemessages: list[dict[str, str]] | None = Noneinput_ids: torch.LongTensor | None = Noneattention_mask: torch.LongTensor | None = Noneuse_chat_template: bool = Trueadd_generation_prompt: bool = Truegen_length: int = 2048block_length: int | None = Nonenum_inference_steps: int = 32temperature: float = 0.0top_p: float | None = Nonetop_k: int | None = Nonesampling_method: str = 'multinomial'threshold: float = 0.7editing_threshold: float | None = 0.5max_post_steps: int = 16minimal_topk: int = 1eos_early_stop: bool = Trueeos_token_id: int | None = Nonemask_token_id: int | None = Nonegenerator: torch.Generator | None = Noneoutput_type: str = 'text'return_dict: bool = Truecallback_on_step_end: Callable[[int, int, dict], None] | PipelineCallback | MultiPipelineCallbacks | None = Nonecallback_on_step_end_tensor_inputs: list[str] | None = None ) → LLaDA2PipelineOutput or tuple

参数

  • 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，可选) — 用于模板的 Mask token ID。
  • generator (torch.Generator，可选) — 用于采样的随机数生成器。
  • output_type (str，默认为 "text") — 输出格式。"text" 将序列解码为字符串（需要 tokenizer）。"seq" 仅返回原始的 token ID 序列。
  • return_dict (bool，可选，默认为 True) — 是否返回 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]，可选) — 传递给回调的张量键。允许的键： block_x、transfer_index、 editing_transfer_index、sampled_tokens、sampled_probs、active_block。

返回值

LLaDA2PipelineOutput 或 tuple

如果 return_dict 为 True，则返回 LLaDA2PipelineOutput；否则返回一个 tuple，其中第一个元素是生成的 token ID（torch.LongTensor），第二个元素是解码后的文本（list[str]）；当 output_type 为 "seq" 时，第二个元素为 None。

使用分块细化生成文本。

示例：

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

class diffusers.LLaDA2PipelineOutput

< >

( sequences: torch.LongTensortexts: list[str] | None = None )

Update on GitHub

←DiffusionGemma Allegro→