Diffusers documentation
BlockRefinementScheduler
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
< source >( 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
< source >( 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.LongTensorde forme(batch_size, seq_len)) — ID de jetons propres. - attention_mask (
torch.LongTensorde 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
< source >( 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
< source >( 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.LongTensorde forme(batch_size, seq_len)) — Séquence complète actuelle incluant tous les blocs jusqu'à la fenêtre courante. - sampled_tokens (
torch.LongTensorde forme(batch_size, block_length)) — Jetons échantillonnés par le planificateur à cette étape. - final_transfer (
torch.BoolTensorde forme(batch_size, block_length)) — Masque combiné des positions validées et modifiées. - finished (
torch.BoolTensorde 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.
Répartir uniformément les block_length jetons sur num_inference_steps étapes.
étape
< source >( 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.Tensorde forme(batch_size, block_length, vocab_size)) — Logits bruts du modèle pour le bloc actuel. - timestep (
intoutorch.Tensor) — Index de l'étape actuelle dans le plan de raffinement du bloc. - sample (
torch.LongTensorde 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 estNone, le planificateur fonctionne en mode de corruption uniforme : il suit les positions validées en interne (réinitialisées lorsquetimestep == 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ùTruemarque 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 unBlockRefinementSchedulerOutputou 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
< source >( 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.LongTensorde forme(batch_size, block_length)) — Tokens du bloc mis à jour après l'étape de raffinement actuelle. - transfer_index (
torch.BoolTensorde forme(batch_size, block_length)) — Masque booléen indiquant quels jetons ont été validés (remplissage du masque). - editing_transfer_index (
torch.BoolTensorde forme(batch_size, block_length)) — Masque booléen indiquant quels jetons ont été modifiés (remplacement non masqué). - sampled_tokens (
torch.LongTensorde forme(batch_size, block_length)) — Identifiants des jetons échantillonnés à partir des logits du modèle. - sampled_probs (
torch.Tensorde forme(batch_size, block_length)) — Probabilités des jetons échantillonnés. - pred_logits (
torch.Tensorde 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.