DprModel

class lightning_ir.models.bi_encoders.dpr.DprModel(config: SingleVectorBiEncoderConfig, *args, **kwargs)

Bases: SingleVectorBiEncoderModel

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

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

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

config_class

Configuration class for a DPR model.

alias of DprConfig

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