Hugging Face's logo

Diffusers documentation

LLaDA2

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

LLaDA2

LLaDA2 est une famille de modèles de langage de diffusion discrète qui génèrent du texte par raffinement itératif par blocs. Plutôt que d’utiliser une génération auto-régressive jeton par jeton, LLaDA2 démarre avec une séquence entièrement masquée et démasque progressivement les jetons en fonction de leur confiance au fil de plusieurs étapes de raffinement.

Utilisation

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

Les callbacks s'exécutent après chaque étape de raffinement. Passez callback_on_step_end_tensor_inputs pour sélectionner les tenseurs inclus dans callback_kwargs. Dans l'implémentation actuelle, block_x (la fenêtre de séquence en cours de raffinement) et transfer_index (masque de commit de remplissage de masque) sont fournis ; renvoyez {"block_x": ...} depuis le callback pour remplacer la fenêtre.

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

Paramètres recommandés

Les modèles LLaDA2.1 prennent en charge deux modes :

Modethresholdediting_thresholdmax_post_steps
Qualité0.70.516
Vitesse0.5None16

Passez editing_threshold=None, 0.0 ou une valeur négative pour désactiver l'édition post-masque.

Pour les modèles LLaDA2.0, désactivez l'édition en passant editing_threshold=None ou 0.0.

Pour tous les modèles : block_length=32, temperature=0.0, num_inference_steps=32.

LLaDA2Pipeline

classe diffusers.LLaDA2Pipeline

< >

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

Pipeline pour la génération de texte par diffusion discrète de type LLaDA2 via raffinement itératif par blocs.

Ce pipeline maintient une séquence de modèle remplie de mask_token_id et l'affine par blocs. À chaque étape d'affinage, il échantillonne des jetons candidats pour le bloc actif et en valide un sous-ensemble en fonction de la confiance.

Le modèle est censé accepter un masque d'attention et position_ids, et renvoyer des logits de forme [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

Paramètres

  • prompt (str ou List[str], facultatif) — Texte de la requête. Lorsque use_chat_template est True (par défaut) et qu'un tokenizer avec un modèle de chat est disponible, la requête est encapsulée dans un message de chat avant la tokenisation.
  • messages (List[Dict[str, str]], facultatif) — Messages de chat à encoder (par ex. [{"role": "user", "content": "Hello"}]). Prend le pas sur prompt lorsqu'il est fourni. Nécessite un tokenizer doté de apply_chat_template.
  • input_ids (torch.LongTensor, facultatif) — Identifiants d’entrée pré-tokenisés. Prend le pas sur prompt et messages.
  • attention_mask (torch.LongTensor, facultatif) — Masque par jeton (1 pour les jetons de prompt valides, 0 pour le bourrage) correspondant à la forme de input_ids. Utilisé uniquement lorsque input_ids est fourni. S’il est omis (et que input_ids est donné), toutes les positions sont considérées comme valides. Lors de la construction des entrées à partir de prompt / messages, le masque du tokeniseur est transmis automatiquement.
  • use_chat_template (bool, par défaut True) — Indique s'il faut envelopper le prompt dans un modèle de chat.
  • add_generation_prompt (bool, par défaut True) — Indique s'il faut ajouter le prompt de génération lors de l'utilisation de modèles de discussion.
  • gen_length (int) — Nombre de jetons à générer.
  • block_length (int, facultatif) — Taille de bloc pour l’affinage. Si non fourni, la valeur de block_length configurée sur le planificateur est utilisée.
  • num_inference_steps (int) — Nombre d'étapes de raffinement par bloc.
  • temperature (float) — Température d'échantillonnage.
  • top_p (float, facultatif) — Seuil d'échantillonnage par noyau.
  • top_k (int, facultatif) — Seuil d'échantillonnage top-k.
  • sampling_method (str) — Méthode d'échantillonnage (auto, greedy, multinomial).
  • threshold (float) — Seuil de confiance pour valider les jetons.
  • editing_threshold (float, facultatif) — Seuil de confiance pour l'édition des jetons déjà validés (hors masque). S'il est positif, une fois que tous les jetons masque d'un bloc sont résolus, le pipeline continue de raffiner : si le modèle prédit un jeton différent avec une confiance supérieure à ce seuil, le jeton existant est remplacé. Définissez-le sur None, 0.0 ou une valeur négative pour désactiver l'édition. La valeur par défaut est 0.5.
  • max_post_steps (int) — Nombre maximal d'itérations de raffinement supplémentaires après la résolution de tous les jetons de masque d'un bloc. Utilisé uniquement lorsque editing_threshold est activé. Valeur par défaut : 16.
  • minimal_topk (int) — Nombre minimal de jetons à valider à chaque étape.
  • eos_early_stop (bool) — Indique s'il faut arrêter après l'engagement de l'EOS dans un bloc.
  • eos_token_id (int, facultatif) — ID du jeton EOS à utiliser pour l’arrêt anticipé.
  • mask_token_id (int, facultatif) — ID du jeton de masque à utiliser pour le modèle.
  • generator (torch.Generator, facultatif) — Générateur aléatoire pour l'échantillonnage.
  • output_type (str, par défaut "text") — Format de sortie. "text" décode les séquences en chaînes de caractères (nécessite un tokenizer). "seq" renvoie uniquement les séquences brutes d'identifiants de jetons.
  • return_dict (bool, facultatif, par défaut True) — Indique s'il faut retourner un LLaDA2PipelineOutput plutôt qu'un tuple.
  • callback_on_step_end (Callable ou PipelineCallback, facultatif) — Callback exécuté après chaque étape de raffinement avec la signature callback_on_step_end(self, step: int, timestep: int, callback_kwargs: Dict).
  • callback_on_step_end_tensor_inputs (List[str], facultatif) — Clés de tenseurs à transmettre au rappel. Clés autorisées : block_x, transfer_index, editing_transfer_index, sampled_tokens, sampled_probs, active_block.

Retourne

LLaDA2PipelineOutput ou tuple

Si return_dict est True, LLaDA2PipelineOutput est retourné, sinon un tuple est retourné dont le premier élément est les identifiants de jetons générés (torch.LongTensor) et le deuxième élément est les textes décodés (list[str]), ou None lorsque output_type est "seq".

Générer du texte avec une raffinement par blocs.

Exemples :

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

< >

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

Update on GitHub

←DiffusionGemma Allegro→