Pooler

class lightning_ir.modeling_utils.embedding_post_processing.Pooler(config: LightningIRConfig)

Bases: Module

Applies pooling to the embeddings based on the pooling strategy defined in the configuration.

__init__(config: LightningIRConfig) → None

Initializes the pooler.

Parameters:

config (LightningIRConfig) – Configuration containing the pooling strategy to apply

Methods

__init__(config)	Initializes the pooler.
forward(embeddings[, attention_mask])	Applies optional pooling to the embeddings.

Attributes

training

forward(embeddings: Tensor, attention_mask: Tensor | None = None) → Tensor

Applies optional pooling to the embeddings.

Parameters:

• embeddings (torch.Tensor) – Query, document, or joint query-document embeddings

• attention_mask (torch.Tensor | None) – Query, document, or joint query-document attention mask

Returns:

(Optionally) pooled embeddings

Return type:

torch.Tensor

Raises:

ValueError – If an unknown pooling strategy is passed