CoilEmbedding

class lightning_ir.models.coil.CoilEmbedding(embeddings: Tensor, scoring_mask: Tensor | None = None, encoding: BatchEncoding | None = None, ids: List[str] | None = None, token_embeddings: Tensor | None = None, cls_embeddings: Tensor | None = None)

Bases: BiEncoderEmbedding

Dataclass containing embeddings and the encoding for COIL models.

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

Methods

__init__(embeddings[, scoring_mask, ...])

Attributes

cls_embeddings	Separate [CLS] token embeddings.
token_embeddings	Token embeddings of the COIL model.
embeddings	Raw embeddings of the COIL model.

cls_embeddings: Tensor | None = None

Separate [CLS] token embeddings.

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

Raw embeddings of the COIL model. Should not be used directly for scoring.

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: torch.Tensor | None = None

Mask tensor designating which vectors should be ignored during scoring.

to(device) → 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

token_embeddings: Tensor | None = None

Token embeddings of the COIL model.