Diffusers documentation
BlockRefinementScheduler
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
< source >( 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
< 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]
Parâmetros
- original_samples (
torch.LongTensorde forma(batch_size, seq_len)) — IDs de tokens limpos. - attention_mask (
torch.LongTensorde 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
< source >( 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
< source >( 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.LongTensorde forma(batch_size, seq_len)) — Sequência completa atual, incluindo todos os blocos até a janela atual. - sampled_tokens (
torch.LongTensorde forma(batch_size, block_length)) — Tokens amostrados pelo agendador nesta etapa. - final_transfer (
torch.BoolTensorde forma(batch_size, block_length)) — Máscara combinada das posições confirmadas e editadas. - finished (
torch.BoolTensorde 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.
Distribua uniformemente os block_length tokens entre as num_inference_steps etapas.
etapa
< 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 )
Parâmetros
- model_output (
torch.Tensorde forma(batch_size, block_length, vocab_size)) — Logits brutos do modelo para o bloco atual. - timestep (
intoutorch.Tensor) — Índice atual da etapa dentro do cronograma de refinamento do bloco. - sample (
torch.LongTensorde 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. QuandoNone, o agendador opera no modo de corrupção uniforme: ele rastreia internamente as posições comprometidas (reiniciando emtimestep == 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,)ondeTruemarca 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 umBlockRefinementSchedulerOutputou uma tupla.
Execute uma etapa única de refinamento: amostrar dos logits, confirmar tokens confiantes e, opcionalmente, editar os existentes.
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 )
Parâmetros
- prev_sample (
torch.LongTensorde forma(batch_size, block_length)) — Tokens do bloco atualizados após a etapa de refinamento atual. - transfer_index (
torch.BoolTensorde forma(batch_size, block_length)) — Máscara booleana que indica quais tokens foram confirmados (preenchimento de máscara). - editing_transfer_index (
torch.BoolTensorde forma(batch_size, block_length)) — Máscara booleana indicando quais tokens foram editados (substituição não-máscara). - sampled_tokens (
torch.LongTensorde forma(batch_size, block_length)) — IDs dos tokens amostrados a partir dos logits do modelo. - sampled_probs (
torch.Tensorde forma(batch_size, block_length)) — Probabilidades dos tokens amostrados. - pred_logits (
torch.Tensorde 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.