diffusers/main/vi/api/schedulers/block_refinement.md

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. Ở 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[[diffusers.BlockRefinementScheduler]]

diffusers.BlockRefinementScheduler[[diffusers.BlockRefinementScheduler]]

diffusers.BlockRefinementScheduler(block_length: int = 32, num_inference_steps: int = 32, threshold: float = 0.95, editing_threshold: float | None = None, minimal_topk: int = 1)

Nguồn

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 cho 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[[diffusers.BlockRefinementScheduler.add_noise]]

add_noise(original_samples: torch.LongTensor, attention_mask: torch.LongTensor, prompt_length: int, block_length: int, mask_token_id: int, generator: torch.Generator | None = None)

Nguồn

Tham số:

original_samples (torch.LongTensor có hình dạng (batch_size, seq_len)) : Các ID token sạch.

attention_mask (torch.LongTensor có hình dạng (batch_size, seq_len)) : Mặt nạ đệm (1 cho hợp lệ, 0 cho đệm).

prompt_length (int) : Số lượng token prompt đầu tiên cần giữ không bị che.

block_length (int) : Kích thước khối cho việc che (masking).

mask_token_id (int) : ID 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[[diffusers.BlockRefinementScheduler.check_block_should_continue]]

check_block_should_continue(step_idx: int, masks_remaining: bool, editing_enabled: bool, editing_transfer_index: torch.BoolTensor, post_steps: int, max_post_steps: int, finished: torch.BoolTensor)

Nguồn

Tham số:

step_idx (int) : Chỉ số bước tinh chỉnh hiện tại trong khối này.

masks_remaining (bool) : Cho biết có còn token mask nào trong khối hay không.

editing_enabled (bool) : Chế độ chỉnh sửa có đang hoạt động hay không.

editing_transfer_index (torch.BoolTensor) : Các token nào đã được chỉnh sửa trong bước này.

post_steps (int) : Số bước chỉnh sửa sau khi áp dụng 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 tối đa cho phép sau khi áp dụng mặt nạ.

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[[diffusers.BlockRefinementScheduler.check_eos_finished]]

check_eos_finished(cur_x: torch.LongTensor, sampled_tokens: torch.LongTensor, final_transfer: torch.BoolTensor, finished: torch.BoolTensor, eos_token_id: int, mask_token_id: int, prompt_length: int)

Nguồn

Tham số:

cur_x (torch.LongTensor có 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.LongTensor có hình dạng (batch_size, block_length)) : Các token được trình điều phối lấy mẫu trong bước này.

final_transfer (torch.BoolTensor có shape (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.BoolTensor có hình dạng (batch_size,)) : Các cờ hoàn thành cho từng batch hiện tại.

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.

get_num_transfer_tokens[[diffusers.BlockRefinementScheduler.get_num_transfer_tokens]]

get_num_transfer_tokens(block_length: int, num_inference_steps: int)

Nguồn

Phân bổ đều các cam kết token block_length trên num_inference_steps bước.

bước[[diffusers.BlockRefinementScheduler.step]]

step(model_output: torch.Tensor, timestep: int | torch.Tensor, sample: torch.LongTensor, mask_token_id: int | None = None, temperature: float = 0.0, top_p: float | None = None, top_k: int | None = None, sampling_method: str = 'auto', threshold: float | None = None, editing_threshold: float | None = None, minimal_topk: int | None = None, prompt_mask: torch.BoolTensor | None = None, generator: torch.Generator | None = None, return_dict: bool = True)

Nguồn

Tham số:

model_output (torch.Tensor có hình dạng (batch_size, block_length, vocab_size)) : Logit thô từ mô hình cho khối hiện tại.

timestep (int hoặc torch.Tensor) : Chỉ số bước hiện tại trong lịch trình tinh chỉnh của khối.

sample (torch.LongTensor có 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, optional) : Token ID được sử dụng cho các vị trí bị che. Khi None, 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 khi timestep == 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, tuỳ chọn) : Ngưỡng cắt cho mẫu hóa hạt nhân (nucleus sampling).

top_k (int, tuỳ 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, tuỳ 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, optional) : Ngưỡng độ tin cậy để chỉnh sửa các token không thuộc mặt nạ; phải là số dương để bật chỉnh sửa. Mặc định là giá trị trong cấu hình.

minimal_topk (int, optional) : Số token tối thiểu cần cam kết cho mỗi bước. Mặc định là giá trị cấu hình.

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) : Xác định việc trả về một BlockRefinementSchedulerOutput hay 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[[diffusers.BlockRefinementSchedulerOutput]]

diffusers.BlockRefinementSchedulerOutput[[diffusers.BlockRefinementSchedulerOutput]]

diffusers.BlockRefinementSchedulerOutput(prev_sample: torch.LongTensor, transfer_index: torch.BoolTensor, editing_transfer_index: torch.BoolTensor, sampled_tokens: torch.LongTensor, sampled_probs: torch.Tensor, pred_logits: torch.Tensor)

Nguồn

Tham số:

prev_sample (torch.LongTensor có 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.BoolTensor có shape (batch_size, block_length)) : Mặt nạ Boolean chỉ ra các token đã được cam kết (mask-filling).

editing_transfer_index (torch.BoolTensor có shape (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.LongTensor có 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.Tensor có 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.Tensor có shape (batch_size, block_length, vocab_size)) : Các logits của denoiser, đượ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.