BiEncoderConfig

class lightning_ir.bi_encoder.config.BiEncoderConfig(similarity_function: Literal['cosine', 'dot'] = 'dot', query_expansion: bool = False, attend_to_query_expanded_tokens: bool = False, query_pooling_strategy: Literal['first', 'mean', 'max', 'sum'] | None = 'mean', query_mask_scoring_tokens: Sequence[str] | Literal['punctuation'] | None = None, query_aggregation_function: Literal['sum', 'mean', 'max', 'harmonic_mean'] = 'sum', doc_expansion: bool = False, attend_to_doc_expanded_tokens: bool = False, doc_pooling_strategy: Literal['first', 'mean', 'max', 'sum'] | None = 'mean', doc_mask_scoring_tokens: Sequence[str] | Literal['punctuation'] | None = None, normalize: bool = False, sparsification: Literal['relu', 'relu_log'] | None = None, add_marker_tokens: bool = False, embedding_dim: int = 768, projection: Literal['linear', 'linear_no_bias', 'mlm'] | None = 'linear', **kwargs)[source]

Bases: LightningIRConfig

The configuration class to instantiate a Bi-Encoder model.

__init__(similarity_function: Literal['cosine', 'dot'] = 'dot', query_expansion: bool = False, attend_to_query_expanded_tokens: bool = False, query_pooling_strategy: Literal['first', 'mean', 'max', 'sum'] | None = 'mean', query_mask_scoring_tokens: Sequence[str] | Literal['punctuation'] | None = None, query_aggregation_function: Literal['sum', 'mean', 'max', 'harmonic_mean'] = 'sum', doc_expansion: bool = False, attend_to_doc_expanded_tokens: bool = False, doc_pooling_strategy: Literal['first', 'mean', 'max', 'sum'] | None = 'mean', doc_mask_scoring_tokens: Sequence[str] | Literal['punctuation'] | None = None, normalize: bool = False, sparsification: Literal['relu', 'relu_log'] | None = None, add_marker_tokens: bool = False, embedding_dim: int = 768, projection: Literal['linear', 'linear_no_bias', 'mlm'] | None = 'linear', **kwargs)[source]

Initializes a bi-encoder configuration.

Parameters:

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

  • query_expansion (bool, optional) – Whether to expand queries with mask tokens, defaults to False

  • attend_to_query_expanded_tokens (bool, optional) – Whether to allow query tokens to attend to mask tokens, defaults to False

  • query_pooling_strategy (Literal['first', 'mean', 'max', 'sum'] | None, optional) – Whether and how to pool the query token embeddings, defaults to “mean”

  • query_mask_scoring_tokens (Sequence[str] | Literal['punctuation'] | None, optional) – Whether and which query tokens to ignore during scoring, defaults to None

  • query_aggregation_function (Literal[ 'sum', 'mean', 'max', 'harmonic_mean' ], optional) – How to aggregate similarity scores over query tokens, defaults to “sum”

  • doc_expansion (bool, optional) – Whether to expand documents with mask tokens, defaults to False

  • attend_to_doc_expanded_tokens (bool, optional) – Whether to allow document tokens to attend to mask tokens, defaults to False

  • doc_pooling_strategy (Literal['first', 'mean', 'max', 'sum'] | None, optional) – Whether andhow to pool document token embeddings, defaults to “mean”

  • doc_mask_scoring_tokens (Sequence[str] | Literal['punctuation'] | None, optional) – Whether and which document tokens to ignore during scoring, defaults to None

  • normalize (bool, optional) – Whether to normalize query and document embeddings, defaults to False

  • sparsification (Literal['relu', 'relu_log'] | None, optional) – Whether and which sparsification function to apply, defaults to None

  • add_marker_tokens (bool, optional) – Whether to add extra marker tokens [Q] / [D] to queries / documents, defaults to False

  • embedding_dim (int, optional) – The output embedding dimension, defaults to 768

  • projection (Literal['linear', 'linear_no_bias', 'mlm'] | None, optional) – Whether and how to project the output emeddings, defaults to “linear”

Methods

__init__([similarity_function, ...])

Initializes a bi-encoder configuration.

from_dict(config_dict, *args, **kwargs)

Loads the configuration from a dictionary.

from_pretrained(...)

get_config_dict(...)

save_pretrained(save_directory[, push_to_hub])

to_added_args_dict()

Outputs a dictionary of the added arguments.

to_dict()

to_tokenizer_dict()

Outputs a dictionary of the tokenizer arguments.

Attributes

ADDED_ARGS

Arguments added to the configuration.

TOKENIZER_ARGS

Arguments for the tokenizer.

backbone_model_type

Backbone model type for the configuration.

model_type

Model type for the configuration.
classmethod from_dict(config_dict: Dict[str, Any], *args, **kwargs) → LightningIRConfig

Loads the configuration from a dictionary. Wraps the transformers.PretrainedConfig.from_dict method to return a derived LightningIRConfig class. See LightningIRConfigClassFactory for more details.

Parameters:

config_dict (Dict[str, Any]) – Configuration dictionary

Raises:

ValueError – If the model type does not match the configuration model type

Returns:

Derived LightningIRConfig class

Return type:

LightningIRConfig

to_added_args_dict() → Dict[str, Any]

Outputs a dictionary of the added arguments.

Returns:

Added arguments

Return type:

Dict[str, Any]

to_tokenizer_dict() → Dict[str, Any]

Outputs a dictionary of the tokenizer arguments.

Returns:

Tokenizer arguments

Return type:

Dict[str, Any]