BiEncoderModel

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

Bases: LightningIRModel, ABC

A bi-encoder model that encodes queries and documents separately and computes a relevance score between them. See BiEncoderConfig for configuration options.

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

Initializes a bi-encoder model given a BiEncoderConfig.

Parameters:: config (BiEncoderConfig) – Configuration for the bi-encoder model.
Raises:: ValueError – If the similarity function is not supported.

Methods

`__init__`(config, args, *kwargs)	Initializes a bi-encoder model given a `BiEncoderConfig`.
`compute_similarity`(query_embeddings, ...[, ...])	Computes the similarity score between all query and document embedding vector pairs.
`encode`(encoding, input_type)	Encodes a batched tokenized text sequences and returns the embeddings and scoring mask.
`encode_doc`(encoding)	Encodes tokenized documents.
`encode_query`(encoding)	Encodes tokenized queries.
`forward`(query_encoding, doc_encoding[, num_docs])	Embeds queries and/or documents and computes relevance scores between them if both are provided.
`score`(output[, num_docs])	Compute relevance scores between queries and documents.

Attributes

training

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

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) – Specifies how many documents are passed per query. If a sequence of integers, len(num_docs) 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 bi-encoder model.

alias of BiEncoderConfig

abstractmethod encode(encoding: BatchEncoding, input_type: 'query' | 'doc') → BiEncoderEmbedding[source]

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 for the text sequence.

Return type:

BiEncoderEmbedding

Raises:

NotImplementedError – If the method is not implemented.

encode_doc(encoding: BatchEncoding) → BiEncoderEmbedding[source]

Encodes tokenized documents.

Parameters:: encoding (BatchEncoding) – Tokenizer encodings for the documents.
Returns:: Document embeddings and scoring mask.
Return type:: BiEncoderEmbedding

encode_query(encoding: BatchEncoding) → BiEncoderEmbedding[source]

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[source]

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. Defaults to None.
doc_encoding (BatchEncoding | None) – Tokenizer encodings for the documents. Defaults to None.
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 of the model containing query and document embeddings and relevance scores.

Return type:

BiEncoderOutput

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

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:

Relevance scores.

Return type:

torch.Tensor

Raises:

NotImplementedError – If the method is not implemented.