Gom cụm Vector Nhúng Chuỗi bằng Cơ Chế Pooling có Mặt Nạ (Masked Sequence Embedding Pooling)

Mô tả bài toán

Khi xây dựng các mô hình Sentence Embedding (như BERT, SimCSE, OpenAI text-embedding), biểu diễn của cả câu được tạo ra bằng cách gom cụm (pooling) các vector token dọc theo trục độ dài chuỗi T.

Do các câu trong một batch có độ dài khác nhau và được đệm bằng <pad>, việc pooling phải tôn trọng attention_mask (giá trị 1 cho token thật, 0 cho token đệm).

Yêu cầu kỹ thuật:

Viết hàm:

def pool_sequence_embeddings(
    embeddings: np.ndarray,
    attention_mask: np.ndarray,
    mode: str = "mean",
    eps: float = 1e-9
) -> np.ndarray

• Đầu vào:
• embeddings: tensor float shape (B, T, D) đại diện cho vector nhúng của B câu.
• attention_mask: tensor số nguyên/nhị phân shape (B, T), trong đó 1 là token thật, 0 là token padding.
• mode: chuỗi chỉ định phép pooling, gồm "mean", "sum", hoặc "max". Nếu nhận giá trị khác, ném ValueError(f"Unsupported pooling mode: {mode}").
• eps: hằng số tránh chia cho 0 khi tính mean trên câu rỗng.

• Quy tắc tính toán:
• "sum": Tính tổng các vector token thật: ∑_t=1^T x_{b, t} × m_{b, t}, output shape (B, D).
• "mean": Tính trung bình chỉ trên các token thật:

x_b = ∑_t=1^T x_{b, t} × m_{b, t}max(∑_t=1^T m_{b, t}, ε)

• "max": Lấy giá trị lớn nhất trên các token thật. Đối với các vị trí m_{b, t} == 0, phải thay thế bằng -10⁹ trước khi lấy np.max(..., axis=1) để đảm bảo token đệm không bao giờ được chọn làm max (ngay cả khi embedding của token thật toàn số âm).

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