BiEncoderEmbedding

class lightning_ir.bi_encoder.bi_encoder_model.BiEncoderEmbedding(embeddings: Tensor, scoring_mask: Tensor | None = None, encoding: BatchEncoding | None = None, ids: List[str] | None = None)

Bases: object

Dataclass containing embeddings and the encoding for bi-encoder models.

__init__(embeddings: Tensor, scoring_mask: Tensor | None = None, encoding: BatchEncoding | None = None, ids: List[str] | None = None) → None

Methods

__init__(embeddings[, scoring_mask, ...])
items()	Iterates over the embeddings attributes and their values like dict.items().
to()	Moves the embeddings to the specified device.

Attributes

device	Returns the device of the embeddings.
encoding	Tokenizer encodings used to generate the embeddings.
ids	List of ids for the embeddings, e.g., query or document ids.
scoring_mask	Mask tensor designating which vectors should be ignored during scoring.
embeddings	Embedding tensor generated by a bi-encoder model of shape [batch_size x seq_len x hidden_size].

property device: device

Returns the device of the embeddings.

Returns:

Device of the embeddings.

Return type:

torch.device

Raises:

ValueError – If the embeddings and scoring_mask are not on the same device.

embeddings: Tensor

Embedding tensor generated by a bi-encoder model of shape [batch_size x seq_len x hidden_size]. The sequence length varies depending on the pooling strategy and the hidden size varies depending on the projection settings.

encoding: BatchEncoding | None = None

Tokenizer encodings used to generate the embeddings.

ids: List[str] | None = None

List of ids for the embeddings, e.g., query or document ids.

items() → Iterable[Tuple[str, Tensor]]

Iterates over the embeddings attributes and their values like dict.items().

Yields:

Tuple[str, torch.Tensor] – Tuple of attribute name and its value.

scoring_mask: Tensor | None = None

Mask tensor designating which vectors should be ignored during scoring.

to(device: device, /) → Self

to(other: Self, /) → Self

Moves the embeddings to the specified device.

Parameters:

device (torch.device | BiEncoderEmbedding) – Device to move the embeddings to or another instance to move to the same device.

Returns:

The instance with embeddings moved to the specified device.

Return type:

Self