MonoModel

class lightning_ir.models.mono.MonoModel(config: MonoConfig, *args, **kwargs)

Bases: CrossEncoderModel

__init__(config: MonoConfig, *args, **kwargs)

A cross-encoder model that jointly encodes a query and document(s). The contextualized embeddings are aggragated into a single vector and fed to a linear layer which computes a final relevance score.

Parameters:

config (MonoConfig) – Configuration for the mono cross-encoder model.

Methods

__init__(config, *args, **kwargs)	A cross-encoder model that jointly encodes a query and document(s).
forward(encoding)	Computes contextualized embeddings for the joint query-document input sequence and computes a relevance score.

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.

config_class

Configuration class for mono cross-encoder models.

alias of MonoConfig

forward(encoding: BatchEncoding) → CrossEncoderOutput

Computes contextualized embeddings for the joint query-document input sequence and computes a relevance score.

Parameters:

encoding (BatchEncoding) – Tokenizer encodings for the joint query-document input sequence.

Returns:

Output of the model.

Return type:

CrossEncoderOutput

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.

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.