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 es una familia de modelos de lenguaje de difusión discreta que generan texto mediante refinamiento iterativo por bloques. En lugar de la generación autoregresiva token a token, LLaDA2 comienza con una secuencia completamente enmascarada y desmascara progresivamente los tokens según su confianza a lo largo de múltiples pasos de refinamiento.

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

Los callbacks se ejecutan después de cada paso de refinamiento. Pasa callback_on_step_end_tensor_inputs para seleccionar qué tensores se incluyen en callback_kwargs. En la implementación actual, se proporcionan block_x (la ventana de secuencia que se está refinando) y transfer_index (máscara de compromiso para rellenar máscaras); devuelve {"block_x": ...} desde el callback para reemplazar la ventana.

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

Los modelos LLaDA2.1 admiten dos modos:

Modothresholdediting_thresholdmax_post_steps
Calidad0.70.516
Velocidad0.5None16

Pasa editing_threshold=None, 0.0 o un valor negativo para desactivar la edición posterior a la máscara.

Para los modelos LLaDA2.0, desactive la edición pasando editing_threshold=None o 0.0.

Para todos los modelos: block_length=32, temperature=0.0, num_inference_steps=32.

LLaDA2Pipeline

clase diffusers.LLaDA2Pipeline

< >

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

Pipeline para la generación de texto mediante difusión discreta de estilo LLaDA2, a través de refinamiento iterativo por bloques.

Este pipeline mantiene una secuencia de plantilla rellena con un mask_token_id y la refina en bloques. En cada paso de refinamiento, muestrea tokens candidatos para el bloque activo y confirma un subconjunto en función de la confianza.

Se espera que el modelo acepte una máscara de atención y position_ids, y que devuelva logits de forma [batch, seq, vocab_size].

__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

Parámetros

  • prompt (str o List[str], opcional) — Texto del prompt. Cuando use_chat_template es True (valor predeterminado) y hay disponible un tokenizador con una plantilla de chat, el texto se envuelve en un mensaje de chat antes de la tokenización.
  • messages (List[Dict[str, str]], opcional) — Mensajes de chat para codificar (p. ej. [{"role": "user", "content": "Hello"}]). Tiene prioridad sobre prompt cuando se proporciona. Requiere un tokenizador con apply_chat_template.
  • input_ids (torch.LongTensor, opcional) — IDs de entrada pretokenizados. Tiene prioridad sobre prompt y messages.
  • attention_mask (torch.LongTensor, opcional) — Máscara por token (1 para tokens de prompt válidos, 0 para relleno) que coincide con la forma de input_ids. Solo se utiliza cuando se proporciona input_ids. Si se omite (y se da input_ids), todas las posiciones se tratan como válidas. Al construir las entradas a partir de prompt / messages, la máscara del tokenizador se transmite automáticamente.
  • use_chat_template (bool, por defecto True) — Si se debe envolver el prompt en una plantilla de chat.
  • add_generation_prompt (bool, por defecto True) — Si se debe añadir el prompt de generación al usar plantillas de chat.
  • gen_length (int) — Número de tokens a generar.
  • block_length (int, opcional) — Tamaño de bloque para el refinamiento. Si no se proporciona, se utiliza el block_length configurado en el programador.
  • num_inference_steps (int) — Número de pasos de refinamiento por bloque.
  • temperature (float) — Temperatura de muestreo.
  • top_p (float, opcional) — Corte de muestreo de núcleo.
  • top_k (int, opcional) — Corte de muestreo top-k.
  • sampling_method (str) — Método de muestreo (auto, greedy, multinomial).
  • threshold (float) — Umbral de confianza para confirmar los tokens.
  • editing_threshold (float, opcional) — Umbral de confianza para editar tokens ya comprometidos (no máscara). Cuando es positivo, una vez que todos los tokens de máscara de un bloque se han resuelto, el pipeline continúa refinando: si el modelo predice un token diferente con una confianza superior a este umbral, el token existente se reemplaza. Establezca None, 0.0 o un valor negativo para deshabilitar la edición. El valor predeterminado es 0.5.
  • max_post_steps (int) — Número máximo de iteraciones adicionales de refinamiento después de que todos los tokens de máscara de un bloque se hayan resuelto. Solo se utiliza cuando editing_threshold está habilitado. El valor predeterminado es 16.
  • minimal_topk (int) — Número mínimo de tokens que se deben confirmar en cada paso.
  • eos_early_stop (bool) — Si se debe detener después de confirmar el EOS en un bloque.
  • eos_token_id (int, opcional) — ID del token EOS para usar en la detención anticipada.
  • mask_token_id (int, opcional) — ID del token de máscara a utilizar para la plantilla.
  • generator (torch.Generator, opcional) — Generador de números aleatorios para el muestreo.
  • output_type (str, por defecto "text") — Formato de salida. "text" decodifica las secuencias en cadenas de texto (requiere un tokenizador). "seq" devuelve únicamente secuencias de ID de tokens.
  • return_dict (bool, opcional, valor predeterminado True) — Si se debe devolver un LLaDA2PipelineOutput en lugar de una tupla.
  • callback_on_step_end (Callable o PipelineCallback, opcional) — Callback ejecutado después de cada paso de refinamiento con la firma callback_on_step_end(self, step: int, timestep: int, callback_kwargs: Dict).
  • callback_on_step_end_tensor_inputs (List[str], opcional) — Claves de tensor para pasar al callback. Claves permitidas: block_x, transfer_index, editing_transfer_index, sampled_tokens, sampled_probs, active_block.

Devuelve

LLaDA2PipelineOutput o tuple

Si return_dict es True, se devuelve LLaDA2PipelineOutput, en caso contrario se devuelve una tuple donde el primer elemento son los IDs de tokens generados (torch.LongTensor) y el segundo elemento son los textos decodificados (list[str]), o None cuando output_type es "seq".

Generar texto con refinamiento por bloques.

Ejemplos:

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

clase diffusers.LLaDA2PipelineOutput

< >

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

Update on GitHub

←DiffusionGemma Allegro→