RankCallback

class lightning_ir.callbacks.callbacks.RankCallback(save_dir: Path | str | None = None, run_name: str | None = None, overwrite: bool = False)

Bases: Callback, _GatherMixin, _OverwriteMixin

__init__(save_dir: Path | str | None = None, run_name: str | None = None, overwrite: bool = False) → None

Callback to write run file of ranked documents to disk.

Parameters:

• save_dir (Path | str | None, optional) – Directory to save run files to. If None, run files will be saved in the models’ directory, defaults to None

• run_name (str | None, optional) – Name of the run file. If None, the dataset’s dataset_id or file name will be used, defaults to None

• overwrite (bool, optional) – Whether to skip or overwrite already existing run files, defaults to False

Methods

__init__([save_dir, run_name, overwrite])	Callback to write run file of ranked documents to disk.
on_test_batch_end(trainer, pl_module, ...[, ...])	Hook to aggregate and write ranking to file.
setup(trainer, pl_module, stage)	Hook to setup the callback.
teardown(trainer, pl_module, stage)	Hook to cleanup the callback.

Attributes

on_test_batch_end(trainer: Trainer, pl_module: LightningIRModule, outputs: LightningIROutput, batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Hook to aggregate and write ranking to file.

Parameters:

• trainer (Trainer) – PyTorch Lightning Trainer

• pl_module (LightningIRModule) – LightningIR Module

• outputs (LightningIROutput) – Scored query documents pairs

• batch (Any) – Batch of input data

• batch_idx (int) – Index of batch in the current dataset

• dataloader_idx (int, optional) – Index of the dataloader, defaults to 0

setup(trainer: Trainer, pl_module: LightningIRModule, stage: str) → None

Hook to setup the callback.

Parameters:

• trainer (Trainer) – PyTorch Lightning Trainer

• pl_module (LightningIRModule) – LightningIR module

• stage (str) – Stage of the trainer, must be “test”

Raises:

• ValueError – If the stage is not “test”

• ValueError – If no save_dir is provided and model_name_or_path is not a path (the model is not local)

teardown(trainer: Trainer, pl_module: BiEncoderModule, stage: str) → None

Hook to cleanup the callback.

Parameters:

• trainer (Trainer) – PyTorch Lightning Trainer

• pl_module (BiEncoderModule) – LightningIR bi-encoder module used for indexing

• stage (str) – Stage of the trainer, must be “test”