BiEncoderEmbedding

class lightning_ir.bi_encoder.model.BiEncoderEmbedding(embeddings: Tensor, scoring_mask: Tensor, encoding: BatchEncoding | None)

Bases: object

Dataclass containing embeddings and scoring mask for bi-encoder models.

__init__(embeddings: Tensor, scoring_mask: Tensor, encoding: BatchEncoding | None) → None

Methods

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

Attributes

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

property device: device

Returns the device of the embeddings.

Raises:

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

Returns:

The device of the embeddings

Return type:

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

Tokenizer encodings used to generate the embeddings.

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

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

Yield:

Tuple of attribute name and its value

Return type:

Iterator[Iterable[Tuple[str, torch.Tensor]]]

scoring_mask: Tensor

Mask tensor designating which vectors should be ignored during scoring.

to(device: device, /) → BiEncoderEmbedding

to(other: BiEncoderEmbedding, /) → BiEncoderEmbedding

Moves the embeddings and scoring mask to the specified device.

Parameters:

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

Returns:

Self

Return type:

BiEncoderEmbedding