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

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}, epsilon)

• "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).

Input

• Hàm pool_sequence_embeddings(embeddings, attention_mask, mode, eps): Các tham số đầu vào chứa dữ liệu Tensor/mảng NumPy hoặc giá trị siêu tham số tương ứng.

Output

• Hàm pool_sequence_embeddings: Trả về kết quả kiểu np.ndarray theo đúng đặc tả kỹ thuật và kích thước quy định.

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

Ví dụ 1