SingleVectorBiEncoderModel

class lightning_ir.bi_encoder.bi_encoder_model.SingleVectorBiEncoderModel(config: SingleVectorBiEncoderConfig, *args, **kwargs)[source]

Bases: BiEncoderModel

A bi-encoder model that encodes queries and documents separately, pools the contextualized embeddings into a single vector, and computes a relevance score based on the similarities between the two vectors. See SingleVectorBiEncoderConfig for configuration options.

__init__(config: SingleVectorBiEncoderConfig, *args, **kwargs) → None[source]

Initializes a single-vector bi-encoder model given a SingleVectorBiEncoderConfig.

Parameters:

config (SingleVectorBiEncoderConfig) – Configuration for the single-vector bi-encoder model

Methods

__init__(config, *args, **kwargs)

Initializes a single-vector bi-encoder model given a SingleVectorBiEncoderConfig.

score(query_embeddings, doc_embeddings[, ...])

Compute relevance scores between queries and documents.

Attributes

training

ALLOW_SUB_BATCHING = True

Flag to allow mini batches of documents for a single query. Set to false for listwise models to ensure correctness.

compute_similarity(query_embeddings: BiEncoderEmbedding, doc_embeddings: BiEncoderEmbedding, num_docs: Sequence[int] | int | None = None) → Tensor

Computes the similarity score between all query and document embedding vector pairs.

Parameters:

  • query_embeddings (BiEncoderEmbedding) – Embeddings of the queries

  • doc_embeddings (BiEncoderEmbedding) – Embeddings of the documents

  • num_docs (Sequence[int] | int | None, optional) – Specifies how many documents are passed per query. If a sequence of integers, len(num_doc) should be equal to the number of queries and sum(num_docs) equal to the number of documents, i.e., the sequence contains one value per query specifying the number of documents for that query. If an integer, assumes an equal number of documents per query. If None, tries to infer the number of documents by dividing the number of documents by the number of queries, defaults to None

Returns:

Similarity scores between all query and document embedding vector pairs

Return type:

torch.Tensor

config_class

Configuration class for the single-vector bi-encoder model.

alias of SingleVectorBiEncoderConfig

abstract encode(encoding: BatchEncoding, input_type: Literal['query', 'doc']) → BiEncoderEmbedding

Encodes a batched tokenized text sequences and returns the embeddings and scoring mask.

Parameters:

  • encoding (BatchEncoding) – Tokenizer encodings for the text sequence

  • input_type (Literal["query", "doc"]) – Type of input, either “query” or “doc”

Returns:

Embeddings and scoring mask

Return type:

BiEncoderEmbedding

encode_doc(encoding: BatchEncoding) → BiEncoderEmbedding

Encodes tokenized documents.

Parameters:

encoding (BatchEncoding) – Tokenizer encodings for the documents

Returns:

Query embeddings and scoring mask

Return type:

BiEncoderEmbedding

encode_query(encoding: BatchEncoding) → BiEncoderEmbedding

Encodes tokenized queries.

Parameters:

encoding (BatchEncoding) – Tokenizer encodings for the queries

Returns:

Query embeddings and scoring mask

Return type:

BiEncoderEmbedding

forward(query_encoding: BatchEncoding | None, doc_encoding: BatchEncoding | None, num_docs: Sequence[int] | int | None = None) → BiEncoderOutput

Embeds queries and/or documents and computes relevance scores between them if both are provided.

Parameters:

  • query_encoding (BatchEncoding | None) – Tokenizer encodings for the queries

  • doc_encoding (BatchEncoding | None) – Tokenizer encodings for the documents

  • num_docs (Sequence[int] | int | None, optional) – Specifies how many documents are passed per query. If a sequence of integers, len(num_doc) should be equal to the number of queries and sum(num_docs) equal to the number of documents, i.e., the sequence contains one value per query specifying the number of documents for that query. If an integer, assumes an equal number of documents per query. If None, tries to infer the number of documents by dividing the number of documents by the number of queries, defaults to None

Returns:

Output of the model

Return type:

BiEncoderOutput

classmethod from_pretrained(model_name_or_path: str | Path, *args, **kwargs) → Self
Loads a pretrained model. Wraps the transformers.PreTrainedModel.from_pretrained method to return a

derived LightningIRModel. See LightningIRModelClassFactory for more details.

param model_name_or_path:

Name or path of the pretrained model

type model_name_or_path:

str | Path

raises ValueError:

If called on the abstract class LightningIRModel and no config is passed

return:

A derived LightningIRModel consisting of a backbone model and a LightningIRModel mixin

rtype:

LightningIRModel

>>> # Loading using model class and backbone checkpoint
>>> type(CrossEncoderModel.from_pretrained("bert-base-uncased"))
<class 'lightning_ir.base.class_factory.CrossEncoderBertModel'>
>>> # Loading using base class and backbone checkpoint
>>> type(LightningIRModel.from_pretrained("bert-base-uncased", config=CrossEncoderConfig()))
<class 'lightning_ir.base.class_factory.CrossEncoderBertModel'>
pooling(embeddings: Tensor, attention_mask: Tensor | None, pooling_strategy: Literal['first', 'mean', 'max', 'sum'] | None) → Tensor

Helper method to apply pooling to the embeddings.

Parameters:

  • embeddings (torch.Tensor) – Query or document embeddings

  • attention_mask (torch.Tensor | None) – Query or document attention mask

  • pooling_strategy (Literal['first', 'mean', 'max', 'sum'] | None) – The pooling strategy. No pooling is applied if None.

Raises:

ValueError – If an unknown pooling strategy is passed

Returns:

(Optionally) pooled embeddings

Return type:

torch.Tensor

score(query_embeddings: BiEncoderEmbedding, doc_embeddings: BiEncoderEmbedding, num_docs: Sequence[int] | int | None = None) → Tensor[source]

Compute relevance scores between queries and documents.

Parameters:

  • query_embeddings (BiEncoderEmbedding) – Embeddings and scoring mask for the queries

  • doc_embeddings (BiEncoderEmbedding) – Embeddings and scoring mask for the documents

  • num_docs (Sequence[int] | int | None, optional) – Specifies how many documents are passed per query. If a sequence of integers, len(num_doc) should be equal to the number of queries and sum(num_docs) equal to the number of documents, i.e., the sequence contains one value per query specifying the number of documents for that query. If an integer, assumes an equal number of documents per query. If None, tries to infer the number of documents by dividing the number of documents by the number of queries, defaults to None

Returns:

Relevance scores

Return type:

torch.Tensor

sparsification(embeddings: Tensor, sparsification_strategy: Literal['relu', 'relu_log'] | None = None) → Tensor

Helper method to apply sparsification to the embeddings.

Parameters:

  • embeddings (torch.Tensor) – Query or document embeddings

  • sparsification_strategy (Literal['relu', 'relu_log'] | None, optional) – The sparsification strategy. No sparsification is applied if None, defaults to None

Raises:

ValueError – If an unknown sparsification strategy is passed

Returns:

(Optionally) sparsified embeddings

Return type:

torch.Tensor