DprModel

class lightning_ir.models.dpr.DprModel(config: SingleVectorBiEncoderConfig, *args, **kwargs)[source]

Bases: SingleVectorBiEncoderModel

A single-vector DPR model. See DprConfig for configuration options.

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

Initializes a DPR model given a DprConfig.

Parameters:: config (SingleVectorBiEncoderConfig) – Configuration for the DPR model.
Raises:: ValueError – If the embedding dimension is not specified in the configuration.

Methods

`__init__`(config, args, *kwargs)	Initializes a DPR model given a `DprConfig`.
`encode`(encoding, input_type)	Encodes a batched tokenized text sequences and returns the embeddings and scoring mask.

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) – 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 a DPR model.

alias of DprConfig

encode(encoding: BatchEncoding, input_type: Literal['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.

Return type:

BiEncoderEmbedding

encode_doc(encoding: BatchEncoding) → BiEncoderEmbedding

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

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. 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

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.

>>> # 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'>
Args:
model_name_or_path (str | Path): Name or path of the pretrained model.

Raises:
ValueError: If called on the abstract class LightningIRModel and no config is passed.

Returns:
LightningIRModel: A derived LightningIRModel consisting of a backbone model and a LightningIRModel mixin.

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.

Returns:

(Optionally) pooled embeddings.

Return type:

torch.Tensor

Raises:

ValueError – If an unknown pooling strategy is passed.

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.

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) – The sparsification strategy. No sparsification is applied if None. Defaults to None.

Returns:

(Optionally) sparsified embeddings.

Return type:

torch.Tensor

Raises:

ValueError – If an unknown sparsification strategy is passed.