Function _to_binary_mask

text_embeddings_connectors.py:33–37

Convert additive mask to binary [B, seq, 1] and zero out padding.

(encoded, encoded_mask)

def _to_binary_mask(encoded, encoded_mask):
    """Convert additive mask to binary [B, seq, 1] and zero out padding."""
    binary_mask = (encoded_mask < 0.000001).to(torch.int64)
    binary_mask = binary_mask.reshape(encoded.shape[0], encoded.shape[1], 1)
    return encoded * binary_mask, binary_mask


def _rescale_norm(x: torch.Tensor, target_dim: int, source_dim: int) -> torch.Tensor: