Tính phân phối xác suất và Greedy Decoding từ Logits (Numerical Softmax)

Mô tả bài toán

Trong mô hình ngôn ngữ tự hồi quy (Autoregressive Language Model), đầu ra của tầng cuối cùng (LM Head) là một vector logits thô z ∈ R^V với V là kích thước từ điển (vocabulary size).

Để dự đoán token tiếp theo:

• Chuyển đổi z thành phân phối xác suất hợp lệ thông qua Softmax ổn định số học:

z_shift = z - max(z) p_i = exp(z_{shift, i})∑_j=1^V exp(z_{shift, j})

• Chiến lược giải mã tham lam (Greedy Decoding) chọn ra token có xác suất cao nhất:

y^* = argmax_i p_i

Hãy viết hàm predict_next_token_probs(logits: np.ndarray) -> tuple[np.ndarray, int, float]:

• Nhận mảng 1D logits kiểu số thực. Nếu không phải 1D hoặc rỗng, raise ValueError.
• Trả về tuple gồm:
• probs: vector phân phối xác suất 1D numpy array float64 có tổng bằng 1.0.
• top1_id: chỉ số token có xác suất lớn nhất (kiểu int).
• top1_prob: giá trị xác suất lớn nhất đó (kiểu float).

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ả.