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
Returns the device of the embeddings.
List of ids for the embeddings, e.g., query or document ids.
Embedding tensor generated by a bi-encoder model of shape [batch_size x seq_len x hidden_size].
Mask tensor designating which vectors should be ignored during scoring.
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.
- 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]]]
- 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: