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

O BlockRefinementScheduler gerencia o refinamento iterativo por bloco para difusão de tokens discretos. A cada etapa, ele confirma os tokens mais confiantes e, opcionalmente, edita tokens já confirmados quando o modelo prevê um token diferente com alta confiança.

Este scheduler é usado por LLaDA2Pipeline.

BlockRefinementScheduler

classe diffusers.BlockRefinementScheduler

< >

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

Agendador para refinamento iterativo por blocos (commit por confiança).

A cada etapa, o agendador amostra tokens candidatos dos logits do modelo e confirma aqueles com a maior confiança. O número de tokens a serem confirmados por etapa é determinado distribuindo uniformemente o comprimento do bloco entre o número de etapas de refinamento.

Suporta opcionalmente edição: após todos os tokens de máscara serem resolvidos, os tokens podem ser substituídos se o modelo prever um token diferente com confiança acima de um editing_threshold positivo (None, 0.0 ou negativo desativa a edição).

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]

Parâmetros

  • original_samples (torch.LongTensor de forma (batch_size, seq_len)) — IDs de tokens limpos.
  • attention_mask (torch.LongTensor de forma (batch_size, seq_len)) — Máscara de preenchimento (1 para válido, 0 para preenchimento).
  • prompt_length (int) — Número de tokens de prompt iniciais a manter sem máscara.
  • block_length (int) — Tamanho do bloco para mascaramento.
  • mask_token_id (int) — ID do token a ser usado para posições mascaradas.
  • generator (torch.Generator, opcional) — RNG para reprodutibilidade.

Retorna

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

(noisy, noisy_rev, masked, masked_rev) — as duas sequências ruidosas complementares e suas máscaras booleanas correspondentes.

Aplique o processo de encaminhamento (adição de ruído) para mascaramento de blocos semi-autorregressivos.

Para cada bloco após o prompt, uma fração aleatória de tokens válidos (não de preenchimento) é substituída por mask_token_id. Duas visões complementares são retornadas: noisy e noisy_rev, nas quais as posições mascaradas em uma correspondem às posições não mascaradas na outra.

check_block_should_continue

< >

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

Parâmetros

  • step_idx (int) — Índice atual da etapa de refinamento dentro deste bloco.
  • masks_remaining (bool) — Se ainda há tokens de máscara no bloco.
  • editing_enabled (bool) — Se o modo de edição está ativo.
  • editing_transfer_index (torch.BoolTensor) — Quais tokens foram editados nesta etapa.
  • post_steps (int) — Número de etapas de edição pós-máscara realizadas até o momento.
  • max_post_steps (int) — Número máximo de etapas de edição pós-máscara permitidas.
  • finished (torch.BoolTensor) — Flags de finalização por lote (da detecção de EOS).

Retorna

bool

True se o refinamento deve continuar, False para interromper.

Determine se o loop de refinamento interno deve continuar para o bloco atual.

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

Parâmetros

  • cur_x (torch.LongTensor de forma (batch_size, seq_len)) — Sequência completa atual, incluindo todos os blocos até a janela atual.
  • sampled_tokens (torch.LongTensor de forma (batch_size, block_length)) — Tokens amostrados pelo agendador nesta etapa.
  • final_transfer (torch.BoolTensor de forma (batch_size, block_length)) — Máscara combinada das posições confirmadas e editadas.
  • finished (torch.BoolTensor de forma (batch_size,)) — Flags de finalização atuais por lote.
  • eos_token_id (int) — ID do token EOS.
  • mask_token_id (int) — ID do token de máscara.
  • prompt_length (int) — Número de tokens de prompt no início da sequência.

Retorna

torch.BoolTensor

Atualizados os flags de finalização.

Atualize as flags de conclusão por lote quando os tokens EOS forem confirmados.

get_num_transfer_tokens

< >

( block_length: intnum_inference_steps: int )

Distribua uniformemente os block_length tokens entre as num_inference_steps etapas.

etapa

< >

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

Parâmetros

  • model_output (torch.Tensor de forma (batch_size, block_length, vocab_size)) — Logits brutos do modelo para o bloco atual.
  • timestep (int ou torch.Tensor) — Índice atual da etapa dentro do cronograma de refinamento do bloco.
  • sample (torch.LongTensor de forma (batch_size, block_length)) — IDs dos tokens do bloco atual (contém tokens de máscara para posições não confirmadas no modo baseado em máscara).
  • mask_token_id (int, opcional) — ID do token usado para posições mascaradas. Quando None, o agendador opera no modo de corrupção uniforme: ele rastreia internamente as posições comprometidas (reiniciando em timestep == 0) e renasce as não comprometidas com tokens aleatórios uniformes, correspondendo ao amostrador de refinamento de blocos do DiffusionGemma.
  • 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, opcional) — Limiar de confiança para confirmar tokens. O valor padrão é o definido na configuração.
  • editing_threshold (float, opcional) — Limiar de confiança para a edição de tokens não mascarados; deve ser positivo para habilitar a edição. O valor padrão é o valor da configuração.
  • minimal_topk (int, opcional) — Número mínimo de tokens a serem confirmados por etapa. O valor padrão é o definido na configuração.
  • prompt_mask (torch.BoolTensor, opcional) — Máscara booleana de forma (block_length,) onde True marca as posições do prompt (não editáveis).
  • generator (torch.Generator, opcional) — Gera números aleatórios para amostragem.
  • return_dict (bool) — Se deve retornar um BlockRefinementSchedulerOutput ou uma tupla.

Execute uma etapa única de refinamento: amostrar dos logits, confirmar tokens confiantes e, opcionalmente, editar os existentes.

BlockRefinementSchedulerOutput

classe diffusers.BlockRefinementSchedulerOutput

< >

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

Parâmetros

  • prev_sample (torch.LongTensor de forma (batch_size, block_length)) — Tokens do bloco atualizados após a etapa de refinamento atual.
  • transfer_index (torch.BoolTensor de forma (batch_size, block_length)) — Máscara booleana que indica quais tokens foram confirmados (preenchimento de máscara).
  • editing_transfer_index (torch.BoolTensor de forma (batch_size, block_length)) — Máscara booleana indicando quais tokens foram editados (substituição não-máscara).
  • sampled_tokens (torch.LongTensor de forma (batch_size, block_length)) — IDs dos tokens amostrados a partir dos logits do modelo.
  • sampled_probs (torch.Tensor de forma (batch_size, block_length)) — Probabilidades dos tokens amostrados.
  • pred_logits (torch.Tensor de forma (batch_size, block_length, vocab_size)) — Os logits do desruidor, transmitidos para autocondicionar o próximo passo.

Classe de saída para agendamento de refinamento de blocos.

Update on GitHub