Diffusers documentation
BlockRefinementScheduler
BlockRefinementScheduler
BlockRefinementScheduler quản lý quá trình tinh chỉnh lặp lại theo khối cho khuếch tán token rời rạc. Tại mỗi bước, nó xác nhận các token có độ tin cậy cao nhất và tùy chọn chỉnh sửa các token đã được xác nhận trước đó khi mô hình dự đoán một token khác với độ tin cậy cao.
Trình lập lịch này được sử dụng bởi LLaDA2Pipeline.
BlockRefinementScheduler
lớp diffusers.BlockRefinementScheduler
< source >( block_length: int = 32num_inference_steps: int = 32threshold: float = 0.95editing_threshold: float | None = Noneminimal_topk: int = 1 )
Trình lập lịch cho tinh chỉnh lặp theo khối (cam kết theo độ tin cậy).
Tại mỗi bước, trình lập lịch lấy mẫu các token ứng viên từ logits của mô hình và cam kết những token có độ tin cậy cao nhất. Số lượng token cần cam kết ở mỗi bước được xác định bằng cách phân bổ đều độ dài khối trên số bước tinh chỉnh.
Tùy chọn hỗ trợ chỉnh sửa: sau khi tất cả các token mặt nạ được giải quyết, các token có thể được thay thế nếu mô hình dự đoán một token khác với độ tin cậy cao hơn ngưỡng editing_threshold dương (None, 0.0 hoặc giá trị âm sẽ vô hiệu hóa việc chỉnh sửa).
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]
Tham số
- original_samples (
torch.LongTensorcó hình dạng(batch_size, seq_len)) — Các ID token sạch. - attention_mask (
torch.LongTensorcó hình dạng(batch_size, seq_len)) — Mặt nạ đệm (1 cho phần hợp lệ, 0 cho phần đệm). - prompt_length (
int) — Số lượng token prompt đầu tiên được giữ không bị che. - block_length (
int) — Kích thước khối để che (masking). - mask_token_id (
int) — ID của token được sử dụng cho các vị trí bị che. - generator (
torch.Generator, tùy chọn) — RNG để đảm bảo khả năng tái lập.
Trả về
tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]
(noisy, noisy_rev, masked, masked_rev) — hai chuỗi nhiễu bổ sung và các mặt nạ boolean tương ứng của chúng.
Áp dụng quá trình tiến (thêm nhiễu) cho việc che khối bán tự hồi quy.
Với mỗi khối sau prompt, một tỷ lệ ngẫu nhiên các token hợp lệ (không phải padding) sẽ được thay thế bằng mask_token_id. Hai góc nhìn bổ trợ được trả về: noisy và noisy_rev, trong đó các vị trí bị mask ở góc nhìn này chính là các vị trí không bị mask ở góc nhìn kia.
check_block_should_continue
< source >( step_idx: intmasks_remaining: boolediting_enabled: boolediting_transfer_index: torch.BoolTensorpost_steps: intmax_post_steps: intfinished: torch.BoolTensor ) → bool
Tham số
- step_idx (
int) — Chỉ số bước tinh chỉnh hiện tại trong khối này. - masks_remaining (
bool) — Liệu còn bất kỳ token mặt nạ nào trong khối hay không. - editing_enabled (
bool) — Xác định xem chế độ chỉnh sửa có đang được bật hay không. - editing_transfer_index (
torch.BoolTensor) — Các token đã được chỉnh sửa trong bước này. - post_steps (
int) — Số bước chỉnh sửa sau mặt nạ đã thực hiện đến thời điểm hiện tại. - max_post_steps (
int) — Số bước chỉnh sửa sau khi áp dụng mặt nạ tối đa được cho phép. - finished (
torch.BoolTensor) — Cờ hoàn thành cho từng batch (từ việc phát hiện EOS).
Trả về
bool
True nếu cần tiếp tục tinh chỉnh, False để thoát.
Xác định xem vòng lặp tinh chỉnh bên trong có nên tiếp tục cho khối hiện tại hay không.
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
Tham số
- cur_x (
torch.LongTensorcó hình dạng(batch_size, seq_len)) — Toàn bộ trình tự hiện tại bao gồm tất cả các khối cho đến cửa sổ hiện tại. - sampled_tokens (
torch.LongTensorcó hình dạng(batch_size, block_length)) — Các token được bộ lập lịch lấy mẫu trong bước này. - final_transfer (
torch.BoolTensorcó hình dạng(batch_size, block_length)) — Mặt nạ kết hợp của các vị trí đã cam kết và đã chỉnh sửa. - finished (
torch.BoolTensorcó hình dạng(batch_size,)) — Các cờ hoàn thành hiện tại cho từng batch. - eos_token_id (
int) — ID của token EOS. - mask_token_id (
int) — ID của token mask. - prompt_length (
int) — Số lượng token prompt ở đầu chuỗi.
Trả về
torch.BoolTensor
Đã cập nhật các cờ hoàn thành.
Cập nhật cờ hoàn thành cho từng lô khi các token EOS được cam kết.
Phân bổ đều các cam kết token block_length trên num_inference_steps bước.
bước
< 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 )
Tham số
- model_output (
torch.Tensorcó hình dạng(batch_size, block_length, vocab_size)) — Logits thô từ mô hình cho khối hiện tại. - timestep (
inthoặctorch.Tensor) — Chỉ số bước hiện tại trong lịch trình tinh chỉnh của khối. - sample (
torch.LongTensorcó hình dạng(batch_size, block_length)) — Các ID token của khối hiện tại (chứa các token mask cho các vị trí chưa được cam kết trong chế độ dựa trên mask). - mask_token_id (
int, tùy chọn) — ID của token được sử dụng cho các vị trí bị che. KhiNone, trình lập lịch chạy ở chế độ phá hủy đồng đều: nó theo dõi các vị trí đã cam kết bên trong (được đặt lại khitimestep == 0) và thêm nhiễu cho các vị trí chưa cam kết bằng các token ngẫu nhiên đồng đều, khớp với trình lấy mẫu tinh chỉnh khối của DiffusionGemma. - temperature (
float) — Nhiệt độ lấy mẫu. - top_p (
float, tùy chọn) — Ngưỡng cắt của nucleus sampling. - top_k (
int, tùy chọn) — Ngưỡng cắt cho lấy mẫu top-k. - sampling_method (
str) — Phương pháp lấy mẫu (auto,greedy,multinomial). - threshold (
float, tùy chọn) — Ngưỡng độ tin cậy để cam kết các token. Mặc định là giá trị cấu hình. - editing_threshold (
float, tùy chọn) — Ngưỡng độ tin cậy để chỉnh sửa các token không bị mask; phải là số dương để bật chỉnh sửa. Mặc định là giá trị trong config. - minimal_topk (
int, tùy chọn) — Số lượng token tối thiểu cần cam kết ở mỗi bước. Mặc định là giá trị trong config. - prompt_mask (
torch.BoolTensor, tùy chọn) — Mặt nạ boolean có hình dạng(block_length,)trong đóTrueđánh dấu các vị trí prompt (không thể chỉnh sửa). - generator (
torch.Generator, tùy chọn) — RNG để lấy mẫu. - return_dict (
bool) — Có nên trả về mộtBlockRefinementSchedulerOutputhay một tuple.
Thực hiện một bước tinh chỉnh duy nhất: lấy mẫu từ logits, cam kết các token có độ tin cậy cao và tùy chọn chỉnh sửa các token hiện có.
BlockRefinementSchedulerOutput
lớp diffusers.BlockRefinementSchedulerOutput
< source >( prev_sample: torch.LongTensortransfer_index: torch.BoolTensorediting_transfer_index: torch.BoolTensorsampled_tokens: torch.LongTensorsampled_probs: torch.Tensorpred_logits: torch.Tensor )
Tham số
- prev_sample (
torch.LongTensorcó hình dạng(batch_size, block_length)) — Các token khối đã cập nhật sau bước tinh chỉnh hiện tại. - transfer_index (
torch.BoolTensorcó hình dạng(batch_size, block_length)) — Mặt nạ Boolean chỉ ra các token đã được cam kết (điền mặt nạ). - editing_transfer_index (
torch.BoolTensorcó hình dạng(batch_size, block_length)) — Mặt nạ boolean chỉ ra các token đã được chỉnh sửa (thay thế không dùng mặt nạ). - sampled_tokens (
torch.LongTensorcó hình dạng(batch_size, block_length)) — Các ID token được lấy mẫu từ logits của mô hình. - sampled_probs (
torch.Tensorcó hình dạng(batch_size, block_length)) — Xác suất của các token được lấy mẫu. - pred_logits (
torch.Tensorcó hình dạng(batch_size, block_length, vocab_size)) — Logits của bộ khử nhiễu, được truyền qua để tự điều kiện cho bước tiếp theo.
Lớp đầu ra cho lịch trình tinh chỉnh khối.