Diffusers documentation
BlockRefinementScheduler
BlockRefinementScheduler
BlockRefinementScheduler mengelola penyempurnaan iteratif per blok untuk difusi token diskrit. Pada setiap langkah, scheduler ini mengunci token yang paling yakin dan secara opsional mengedit token yang sudah dikunci ketika model memprediksi token lain dengan keyakinan tinggi.
Scheduler ini digunakan oleh LLaDA2Pipeline.
BlockRefinementScheduler
kelas diffusers.BlockRefinementScheduler
< source >( block_length: int = 32num_inference_steps: int = 32threshold: float = 0.95editing_threshold: float | None = Noneminimal_topk: int = 1 )
Penjadwal untuk penyempurnaan iteratif per blok (commit-by-confidence).
Pada setiap langkah, penjadwal mengambil sampel token kandidat dari logit model dan mengkonfirmasi token dengan kepercayaan tertinggi. Jumlah token yang dikonfirmasi per langkah ditentukan dengan mendistribusikan panjang blok secara merata di antara jumlah langkah penyempurnaan.
Secara opsional mendukung pengeditan: setelah semua token mask teresolusi, token dapat diganti jika model memprediksi token yang berbeda dengan kepercayaan di atas editing_threshold positif (None, 0.0, atau nilai negatif menonaktifkan pengeditan).
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]
Parameter
- original_samples (
torch.LongTensordengan bentuk(batch_size, seq_len)) — ID token bersih. - attention_mask (
torch.LongTensordengan bentuk(batch_size, seq_len)) — Mask padding (1 untuk valid, 0 untuk padding). - prompt_length (
int) — Jumlah token prompt awal yang dipertahankan tanpa masker. - block_length (
int) — Ukuran blok untuk masking. - mask_token_id (
int) — ID token yang digunakan untuk posisi yang dimasker. - generator (
torch.Generator, opsional) — RNG untuk reproduktibilitas.
Mengembalikan
tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]
(noisy, noisy_rev, masked, masked_rev) — dua urutan berisik yang saling melengkapi dan mask boolean yang sesuai.
Terapkan proses maju (noising) untuk masking blok semi-otoregresif.
Untuk setiap blok setelah prompt, sebagian acak dari token valid (non-padding) diganti dengan mask_token_id. Dua tampilan komplementer dikembalikan: noisy dan noisy_rev, di mana posisi yang dimasker pada salah satu adalah posisi yang tidak dimasker pada yang lain.
check_block_should_continue
< source >( step_idx: intmasks_remaining: boolediting_enabled: boolediting_transfer_index: torch.BoolTensorpost_steps: intmax_post_steps: intfinished: torch.BoolTensor ) → bool
Parameter
- step_idx (
int) — Indeks langkah penyempurnaan saat ini di dalam blok ini. - masks_remaining (
bool) — Apakah masih ada token mask yang tersisa di blok. - editing_enabled (
bool) — Apakah mode pengeditan sedang aktif. - editing_transfer_index (
torch.BoolTensor) — Token mana yang diedit pada langkah ini. - post_steps (
int) — Jumlah langkah pengeditan pasca-masker yang telah dilakukan sejauh ini. - max_post_steps (
int) — Jumlah langkah pengeditan pasca-masker maksimum yang diizinkan. - finished (
torch.BoolTensor) — Flag selesai per batch (dari deteksi EOS).
Mengembalikan
bool
True jika penyempurnaan harus dilanjutkan, False untuk berhenti.
Tentukan apakah loop penyempurnaan internal harus dilanjutkan untuk blok saat ini.
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
Parameter
- cur_x (
torch.LongTensordengan bentuk(batch_size, seq_len)) — Urutan penuh saat ini yang mencakup semua blok hingga jendela saat ini. - sampled_tokens (
torch.LongTensordengan bentuk(batch_size, block_length)) — Token yang disampelkan oleh scheduler pada langkah ini. - final_transfer (
torch.BoolTensordengan bentuk(batch_size, block_length)) — Mask gabungan dari posisi yang dikomit dan diedit. - finished (
torch.BoolTensordengan bentuk(batch_size,)) — Flag selesai per-batch saat ini. - eos_token_id (
int) — ID token EOS. - mask_token_id (
int) — ID token mask. - prompt_length (
int) — Jumlah token prompt di awal urutan.
Mengembalikan
torch.BoolTensor
Perbarui bendera selesai.
Perbarui flag selesai per-batch ketika token EOS dikomit.
Distribusikan block_length token commits secara merata di antara num_inference_steps langkah.
langkah
< 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 )
Parameter
- model_output (
torch.Tensordengan bentuk(batch_size, block_length, vocab_size)) — Logit mentah dari model untuk blok saat ini. - timestep (
intatautorch.Tensor) — Indeks langkah saat ini dalam jadwal penyempurnaan blok. - sample (
torch.LongTensordengan bentuk(batch_size, block_length)) — ID token blok saat ini (mengandung token mask untuk posisi yang belum dikomit dalam mode berbasis mask). - mask_token_id (
int, opsional) — ID token yang digunakan untuk posisi yang dimasker. KetikaNone, penjadwal berjalan dalam mode korupsi seragam: ia melacak posisi yang sudah dikomit secara internal (di-reset padatimestep == 0) dan menambahkan noise kembali pada posisi yang belum dikomit dengan token acak seragam, sesuai dengan sampler penyempurnaan blok DiffusionGemma. - temperature (
float) — Suhu sampling. - top_p (
float, opsional) — Batas pemotongan nucleus sampling. - top_k (
int, opsional) — Batas pemotongan sampling top-k. - sampling_method (
str) — Metode sampling (auto,greedy,multinomial). - threshold (
float, opsional) — Ambang batas kepercayaan untuk mengunci token. Bawaan ke nilai konfigurasi. - editing_threshold (
float, opsional) — Ambang batas kepercayaan untuk mengedit token non-mask; harus bernilai positif untuk mengaktifkan pengeditan. Secara default menggunakan nilai config. - minimal_topk (
int, opsional) — Token minimum yang harus dikomit per langkah. Bawaan ke nilai konfigurasi. - prompt_mask (
torch.BoolTensor, opsional) — Mask boolean dengan bentuk(block_length,)di manaTruemenandai posisi prompt (tidak dapat diedit). - generator (
torch.Generator, opsional) — RNG untuk sampling. - return_dict (
bool) — Apakah akan mengembalikanBlockRefinementSchedulerOutputatau tuple.
Lakukan satu langkah penyempurnaan: ambil sampel dari logits, tetapkan token yang yakin, dan secara opsional edit yang sudah ada.
BlockRefinementSchedulerOutput
kelas diffusers.BlockRefinementSchedulerOutput
< source >( prev_sample: torch.LongTensortransfer_index: torch.BoolTensorediting_transfer_index: torch.BoolTensorsampled_tokens: torch.LongTensorsampled_probs: torch.Tensorpred_logits: torch.Tensor )
Parameter
- prev_sample (
torch.LongTensordengan bentuk(batch_size, block_length)) — Token blok yang diperbarui setelah langkah penyempurnaan saat ini. - transfer_index (
torch.BoolTensordengan bentuk(batch_size, block_length)) — Mask boolean yang menunjukkan token mana yang telah dikomit (mask-filling). - editing_transfer_index (
torch.BoolTensordengan bentuk(batch_size, block_length)) — Mask boolean yang menunjukkan token mana yang diedit (penggantian non-mask). - sampled_tokens (
torch.LongTensordengan bentuk(batch_size, block_length)) — ID token yang disampelkan dari logits model. - sampled_probs (
torch.Tensordengan bentuk(batch_size, block_length)) — Probabilitas dari token yang disampelkan. - pred_logits (
torch.Tensordengan bentuk(batch_size, block_length, vocab_size)) — Logits denoiser, diteruskan untuk self-conditioning langkah berikutnya.
Kelas output untuk penjadwalan penyempurnaan blok.