CrossEncoderModule

class lightning_ir.cross_encoder.cross_encoder_module.CrossEncoderModule(model_name_or_path: str | None = None, config: CrossEncoderConfig | None = None, model: CrossEncoderModel | None = None, BackboneModel: Type[PreTrainedModel] | None = None, loss_functions: Sequence[LossFunction | Tuple[LossFunction, float]] | None = None, evaluation_metrics: Sequence[str] | None = None, model_kwargs: Mapping[str, Any] | None = None)[source]

Bases: LightningIRModule

__init__(model_name_or_path: str | None = None, config: CrossEncoderConfig | None = None, model: CrossEncoderModel | None = None, BackboneModel: Type[PreTrainedModel] | None = None, loss_functions: Sequence[LossFunction | Tuple[LossFunction, float]] | None = None, evaluation_metrics: Sequence[str] | None = None, model_kwargs: Mapping[str, Any] | None = None)[source]

LightningIRModule for cross-encoder models. It contains a CrossEncoderModel and a CrossEncoderTokenizer and implements the training, validation, and testing steps for the model.

Parameters:

  • model_name_or_path (str | None) – Name or path of backbone model or fine-tuned Lightning IR model. Defaults to None.

  • config (CrossEncoderConfig | None) – CrossEncoderConfig to apply when loading from backbone model. Defaults to None.

  • model (CrossEncoderModel | None) – Already instantiated CrossEncoderModel. Defaults to None.

  • BackboneModel (Type[PreTrainedModel] | None) – Huggingface PreTrainedModel class to use as backbone instead of the default AutoModel. Defaults to None.

  • loss_functions (Sequence[LossFunction | Tuple[LossFunction, float]] | None) – Loss functions to apply during fine-tuning, optional loss weights can be provided per loss function. Defaults to None.

  • evaluation_metrics (Sequence[str] | None) – Metrics corresponding to ir-measures measure strings to apply during validation or testing. Defaults to None.

  • model_kwargs (Mapping[str, Any] | None) – Additional keyword arguments to pass to from_pretrained when loading a model. Defaults to None.

Methods

__init__([model_name_or_path, config, ...])

LightningIRModule for cross-encoder models.

forward(batch)

Runs a forward pass of the model on a batch of data and returns the contextualized embeddings from the backbone model as well as the relevance scores.

Attributes

training

forward(batch: RankBatch | TrainBatch | SearchBatch) → CrossEncoderOutput[source]

Runs a forward pass of the model on a batch of data and returns the contextualized embeddings from the backbone model as well as the relevance scores.

Parameters:

batch (RankBatch | TrainBatch | SearchBatch) – Batch of data to run the forward pass on.

Returns:

Output of the model.

Return type:

CrossEncoderOutput

Raises:

ValueError – If the batch is a SearchBatch.