Giải Thuật Tìm Kiếm Chùm (Beam Search) Có Phạt Độ Dài (Wu et al.)

Mục tiêu kiến thức

• Khi sinh câu văn từ mô hình Sequence-to-Sequence (Dịch máy, Tóm tắt văn bản), giải thuật tham lam (Greedy) chỉ chọn token tốt nhất tại từng bước có thể bỏ lỡ câu tối ưu toàn cục.
• Beam Search: Duy trì một tập hợp gồm B chùm giả thuyết (hypotheses) có log xác suất tích lũy cao nhất tại mỗi bước.
• Hiện tượng thiên vị câu ngắn (Short sentence bias):

Vì mỗi bước cộng thêm một số âm log P(w_t) ≤ 0, câu càng dài thì tổng log xác suất càng âm. Nếu chỉ so sánh điểm thô, mô hình luôn có xu hướng sinh câu cụt ngủn.

• Hệ số phạt độ dài (Length Penalty - Wu et al., Google GNMT 2016):

Score(Y) = ∑_t=1^|Y| log P(y_t | y_<t)LP(|Y|), LP(|Y|) = (5 + |Y|)^α(5 + 1)^α Với α = 0.6 (hoặc 0.7), hệ số LP(|Y|) bù trừ công bằng cho các chuỗi dài.

Yêu cầu

Viết hàm:

def beam_search_step(
    beam: list[tuple[list[int], float]],
    log_probs: torch.Tensor,
    beam_width: int = 3,
    alpha: float = 0.6
) -> list[tuple[list[int], float]]:
    pass

• beam: Danh sách các tuple (token_sequence, cumulative_log_prob) từ bước trước.
• log_probs: Tensor 1D shape (vocab_size,) chứa log xác suất của token tiếp theo (giả lập cho chùm).
• Trả về top beam_width cặp (sequence, cumulative_log_prob) có điểm số chuẩn hóa cao nhất.

Ràng buộc

• Thời gian chạy tối đa: 2000ms.
• Giới hạn bộ nhớ: 256MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.