MvrConfig

class lightning_ir.models.mvr.MvrConfig(query_length: int = 32, doc_length: int = 512, similarity_function: Literal['cosine', 'dot'] = 'dot', normalize: bool = False, add_marker_tokens: bool = False, embedding_dim: int | None = None, projection: Literal['linear', 'linear_no_bias'] | None = 'linear', num_viewer_tokens: int | None = 8, **kwargs)[source]

Bases: MultiVectorBiEncoderConfig

Configuration class for a MVR model.

__init__(query_length: int = 32, doc_length: int = 512, similarity_function: Literal['cosine', 'dot'] = 'dot', normalize: bool = False, add_marker_tokens: bool = False, embedding_dim: int | None = None, projection: Literal['linear', 'linear_no_bias'] | None = 'linear', num_viewer_tokens: int | None = 8, **kwargs)[source]

A MVR model encodes queries and document separately. It uses a single vector to represent the query and multiple vectors to represent the document. The document representation is obtained from n viewer tokens ([VIE]) prepended to the document. During training, a contrastive loss pushes the viewer token representations away from one another, such that they represent different “views” of the document. Only the maximum similarity between the query vector and the viewer token vectors is used to compute the relevance score.

Parameters:

  • query_length (int) – Maximum query length. Defaults to 32.

  • doc_length (int) – Maximum document length. Defaults to 512.

  • similarity_function (Literal['cosine', 'dot']) – Similarity function to compute scores between query and document embeddings. Defaults to “dot”.

  • normalize (bool) – Whether to normalize query and document embeddings. Defaults to False.

  • add_marker_tokens (bool) – Whether to prepend extra marker tokens [Q] / [D] to queries / documents. Defaults to False.

  • embedding_dim (int | None) – Dimension of the final embeddings. If None, it will be set to the hidden size of the backbone model. Defaults to None.

  • projection (Literal["linear", "linear_no_bias"] | None) – Type of projection layer to apply on the pooled embeddings. If None, no projection is applied. Defaults to “linear”.

  • num_viewer_tokens (int | None) – Number of viewer tokens to prepend to the document. Defaults to 8.

Methods

__init__([query_length, doc_length, ...])

A MVR model encodes queries and document separately.

Attributes

model_type

Model type for a MVR model.
backbone_model_type: str | None = None

Backbone model type for the configuration. Set by LightningIRModelClassFactory().

classmethod from_pretrained(pretrained_model_name_or_path: str | Path, *args, **kwargs) → LightningIRConfig

Loads the configuration from a pretrained model. Wraps the transformers.PretrainedConfig.from_pretrained

Parameters:

pretrained_model_name_or_path (str | Path) – Pretrained model name or path.

Returns:

Derived LightningIRConfig class.

Return type:

LightningIRConfig

Raises:

ValueError – If pretrained_model_name_or_path is not a Lightning IR model and no LightningIRConfig is passed.

get_tokenizer_kwargs(Tokenizer: Type[LightningIRTokenizer]) → Dict[str, Any]

Returns the keyword arguments for the tokenizer. This method is used to pass the configuration parameters to the tokenizer.

Parameters:

Tokenizer (Type[LightningIRTokenizer]) – Class of the tokenizer to be used.

Returns:

Keyword arguments for the tokenizer.

Return type:

Dict[str, Any]

model_type: str = 'mvr'

Model type for a MVR model.

to_dict() → Dict[str, Any]

Overrides the transformers.PretrainedConfig.to_dict method to include the added arguments and the backbone model type.

Returns:

Configuration dictionary.

Return type:

Dict[str, Any]

to_diff_dict() → dict[str, Any]

Removes all attributes from the configuration that correspond to the default config attributes for better readability, while always retaining the config attribute from the class. Serializes to a Python dictionary.

Returns:

Dictionary of all the attributes that make up this configuration instance.

Return type:

dict[str, Any]