BiEncoderEmbedding

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

Bases: object

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

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

Methods

__init__(embeddings, scoring_mask, encoding)

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.

ids

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

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.

ids: List[str] | None = None

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

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

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 | None

Mask tensor designating which vectors should be ignored during scoring.

to(device: device, /) Self[source]
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:

Self

Return type:

BiEncoderEmbedding