Function average_pool

document_qa/deployment/_embeddings_app.py:64–67

Mean-pool token embeddings, ignoring padding positions.

(last_hidden_states: Tensor, attention_mask: Tensor)

def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
    """Mean-pool token embeddings, ignoring padding positions."""
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]


def load_embedding_model(model_name: str, model_revision: str):