SingleVectorBiEncoderModel

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

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

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(output[, num_docs])	Compute relevance scores between queries and documents.

Attributes

training

config_class

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

alias of SingleVectorBiEncoderConfig

score(output: BiEncoderOutput, num_docs: Sequence[int] | int | None = None) → BiEncoderOutput

Compute relevance scores between queries and documents.

Parameters:

• output (BiEncoderOutput) – Output containing embeddings and scoring mask.

• num_docs (Sequence[int] | int | None) – 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 containing relevance scores.

Return type:

BiEncoderOutput

Raises:

ValueError – If query or document embeddings are not provided in the output.