Diffusers documentation

BlockRefinementScheduler

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

BlockRefinementScheduler

Le BlockRefinementScheduler gère le raffinement itératif par blocs pour la diffusion de jetons discrets. À chaque étape, il valide les jetons les plus sûrs et modifie éventuellement les jetons déjà validés lorsque le modèle prédit un jeton différent avec une forte confiance.

Ce planificateur est utilisé par LLaDA2Pipeline.

BlockRefinementScheduler

classe diffusers.BlockRefinementScheduler

< >

( block_length: int = 32num_inference_steps: int = 32threshold: float = 0.95editing_threshold: float | None = Noneminimal_topk: int = 1 )

Planificateur pour le raffinement itératif par blocs (engagement par confiance).

À chaque étape, le planificateur échantillonne des jetons candidats à partir des logits du modèle et valide ceux qui présentent la plus forte confiance. Le nombre de jetons à valider par étape est déterminé en répartissant uniformément la longueur du bloc sur le nombre d'étapes de raffinement.

Prend en charge la modification en option : une fois tous les jetons de masque résolus, les jetons peuvent être remplacés si le modèle prédit un jeton différent avec une confiance supérieure à un editing_threshold positif (None, 0.0 ou une valeur négative désactive la modification).

add_noise

< >

( original_samples: torch.LongTensorattention_mask: torch.LongTensorprompt_length: intblock_length: intmask_token_id: intgenerator: torch.Generator | None = None ) → tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]

Paramètres

  • original_samples (torch.LongTensor de forme (batch_size, seq_len)) — ID de jetons propres.
  • attention_mask (torch.LongTensor de forme (batch_size, seq_len)) — Masque de bourrage (1 pour les positions valides, 0 pour le bourrage).
  • prompt_length (int) — Nombre de jetons de prompt initiaux à conserver non masqués.
  • block_length (int) — Taille du bloc pour le masquage.
  • mask_token_id (int) — Identifiant du jeton à utiliser pour les positions masquées.
  • generator (torch.Generator, facultatif) — Générateur aléatoire pour la reproductibilité.

Retourne

tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]

(noisy, noisy_rev, masked, masked_rev) — les deux séquences bruitées complémentaires et leurs masques booléens correspondants.

Appliquer le processus avant (bruitage) pour le masquage de blocs semi-autorégressifs.

Pour chaque bloc après le prompt, une fraction aléatoire de jetons valides (non de bourrage) est remplacée par mask_token_id. Deux vues complémentaires sont renvoyées : noisy et noisy_rev, où les positions masquées dans l'une correspondent aux positions non masquées dans l'autre.

check_block_should_continue

< >

( step_idx: intmasks_remaining: boolediting_enabled: boolediting_transfer_index: torch.BoolTensorpost_steps: intmax_post_steps: intfinished: torch.BoolTensor ) → bool

Paramètres

  • step_idx (int) — Index de l'étape de raffinement actuelle au sein de ce bloc.
  • masks_remaining (bool) — Indique s'il reste des jetons de masque dans le bloc.
  • editing_enabled (bool) — Indique si le mode d'édition est actif.
  • editing_transfer_index (torch.BoolTensor) — Tokens qui ont été modifiés à cette étape.
  • post_steps (int) — Nombre d'étapes d'édition post-masque effectuées jusqu'à présent.
  • max_post_steps (int) — Nombre maximal de étapes d'édition post-masque autorisées.
  • finished (torch.BoolTensor) — Indicateurs de fin par lot (issus de la détection de EOS).

Retourne

bool

True si le raffinement doit se poursuivre, False pour l’interrompre.

Déterminer si la boucle de raffinement interne doit se poursuivre pour le bloc actuel.

check_eos_finished

< >

( cur_x: torch.LongTensorsampled_tokens: torch.LongTensorfinal_transfer: torch.BoolTensorfinished: torch.BoolTensoreos_token_id: intmask_token_id: intprompt_length: int ) → torch.BoolTensor

Paramètres

  • cur_x (torch.LongTensor de forme (batch_size, seq_len)) — Séquence complète actuelle incluant tous les blocs jusqu'à la fenêtre courante.
  • sampled_tokens (torch.LongTensor de forme (batch_size, block_length)) — Jetons échantillonnés par le planificateur à cette étape.
  • final_transfer (torch.BoolTensor de forme (batch_size, block_length)) — Masque combiné des positions validées et modifiées.
  • finished (torch.BoolTensor de forme (batch_size,)) — Drapeaux de fin actuels par lot.
  • eos_token_id (int) — ID du jeton EOS.
  • mask_token_id (int) — ID du jeton de masquage.
  • prompt_length (int) — Nombre de jetons de l'invite au début de la séquence.

Retourne

torch.BoolTensor

Mise à jour des indicateurs de fin.

Mettre à jour les indicateurs de fin par lot lorsque les jetons EOS sont validés.

get_num_transfer_tokens

< >

( block_length: intnum_inference_steps: int )

Répartir uniformément les block_length jetons sur num_inference_steps étapes.

étape

< >

( model_output: torch.Tensortimestep: int | torch.Tensorsample: torch.LongTensormask_token_id: int | None = Nonetemperature: float = 0.0top_p: float | None = Nonetop_k: int | None = Nonesampling_method: str = 'auto'threshold: float | None = Noneediting_threshold: float | None = Noneminimal_topk: int | None = Noneprompt_mask: torch.BoolTensor | None = Nonegenerator: torch.Generator | None = Nonereturn_dict: bool = True )

Paramètres

  • model_output (torch.Tensor de forme (batch_size, block_length, vocab_size)) — Logits bruts du modèle pour le bloc actuel.
  • timestep (int ou torch.Tensor) — Index de l'étape actuelle dans le plan de raffinement du bloc.
  • sample (torch.LongTensor de forme (batch_size, block_length)) — Identifiants des jetons du bloc actuel (contient des jetons de masque pour les positions non validées en mode basé sur le masque).
  • mask_token_id (int, facultatif) — Identifiant du jeton utilisé pour les positions masquées. Lorsque la valeur est None, le planificateur fonctionne en mode de corruption uniforme : il suit les positions validées en interne (réinitialisées lorsque timestep == 0) et réintroduit du bruit sur les positions non validées à l’aide de jetons aléatoires uniformes, conformément à l’échantillonneur de raffinement par blocs de DiffusionGemma.
  • 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, facultatif) — Seuil de confiance pour valider les jetons. Valeur par défaut : valeur de la configuration.
  • editing_threshold (float, facultatif) — Seuil de confiance pour l'édition des jetons non masqués ; doit être positif pour activer l'édition. Valeur par défaut : valeur de configuration.
  • minimal_topk (int, facultatif) — Nombre minimal de jetons à valider à chaque étape. La valeur par défaut est celle de la configuration.
  • prompt_mask (torch.BoolTensor, facultatif) — Masque booléen de forme (block_length,) où True marque les positions du prompt (non modifiables).
  • generator (torch.Generator, facultatif) — Générateur aléatoire pour l'échantillonnage.
  • return_dict (bool) — Indique s'il faut retourner un BlockRefinementSchedulerOutput ou un tuple.

Effectuer une étape unique de raffinement : échantillonner à partir des logits, valider les jetons confiants et, le cas échéant, modifier les jetons existants.

BlockRefinementSchedulerOutput

classe diffusers.BlockRefinementSchedulerOutput

< >

( prev_sample: torch.LongTensortransfer_index: torch.BoolTensorediting_transfer_index: torch.BoolTensorsampled_tokens: torch.LongTensorsampled_probs: torch.Tensorpred_logits: torch.Tensor )

Paramètres

  • prev_sample (torch.LongTensor de forme (batch_size, block_length)) — Tokens du bloc mis à jour après l'étape de raffinement actuelle.
  • transfer_index (torch.BoolTensor de forme (batch_size, block_length)) — Masque booléen indiquant quels jetons ont été validés (remplissage du masque).
  • editing_transfer_index (torch.BoolTensor de forme (batch_size, block_length)) — Masque booléen indiquant quels jetons ont été modifiés (remplacement non masqué).
  • sampled_tokens (torch.LongTensor de forme (batch_size, block_length)) — Identifiants des jetons échantillonnés à partir des logits du modèle.
  • sampled_probs (torch.Tensor de forme (batch_size, block_length)) — Probabilités des jetons échantillonnés.
  • pred_logits (torch.Tensor de forme (batch_size, block_length, vocab_size)) — Les logits du débruiteur, transmis pour l'auto-conditionnement de l'étape suivante.

Classe de sortie pour la planification du raffinement de blocs.

Update on GitHub