Cơ chế lọc Top-K và che mặt nạ Logits trong LLM (Top-K Sampling)

Mô tả bài toán

Trong quá trình sinh văn bản của LLM (như GPT, LLaMA), để ngăn mô hình chọn phải các từ vựng vô nghĩa ở đuôi phân phối (tail of distribution), kỹ thuật Top-K Sampling chỉ giữ lại K token có logit cao nhất và gán toàn bộ các token còn lại bằng -∞: z'_i = z_i & nếu z_i ≥ z_(K)
-∞ & ngược lại Khi đó sau khi đi qua hàm Softmax, xác suất của các token bị loại sẽ bằng đúng 0 (e^-∞ = 0).

Hãy viết hàm top_k_filtering(logits: np.ndarray, top_k: int, filter_value: float = -float('inf')) -> np.ndarray:

• logits: mảng 1D hoặc 2D NumPy float64 (Batch, Vocab) chứa điểm số logit.
• top_k: số nguyên dương 1 ≤ K ≤ V. Nếu K ≥ V, giữ nguyên toàn bộ logits.
• filter_value: giá trị thay thế (mặc định -∞).
• Trả về mảng NumPy cùng kích thước với các token ngoài top-k bị thay bằng filter_value.

Ràng buộc

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