LLaDA2

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

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.

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

classe diffusers.LLaDA2Pipeline

< source >

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

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

< source >

( 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

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 no 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 de 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 editar 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 confirmar o 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 template.
generator (torch.Generator, opcional) — Gera números aleatórios 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 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 ou tuple

Se return_dict for True, LLaDA2PipelineOutput é retornado, caso contrário, uma tuple é retornada, onde 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:

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

classe diffusers.LLaDA2PipelineOutput

< source >

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

Update on GitHub

Diffusers

LLaDA2

Uso

Callbacks

Parâmetros recomendados

LLaDA2Pipeline

classe diffusers.LLaDA2Pipeline

__call__

LLaDA2PipelineOutput

classe diffusers.LLaDA2PipelineOutput

call