Merge pull request #64 from fschlatt/main

Fix pretty printing for skipped inference datasets

Author: fschlatt

lightning_ir/base/module.py

@@ -14,8 +14,7 @@
 from lightning.pytorch.trainer.states import RunningStage
 from transformers import BatchEncoding
 
-from ..data import RankBatch, RunDataset, SearchBatch, TrainBatch
-from ..data.dataset import IRDataset
+from ..data import IRDataset, RankBatch, RunDataset, SearchBatch, TrainBatch
 from ..loss.loss import InBatchLossFunction, LossFunction
 from .config import LightningIRConfig
 from .model import LightningIRModel, LightningIROutput
@@ -235,7 +234,8 @@ def validation_step(
         if self.evaluation_metrics is None:
             return output
 
-        dataset_id = self.get_dataset_id(dataloader_idx)
+        dataset = self.get_dataset(dataloader_idx)
+        dataset_id = str(dataloader_idx) if dataset is None else self.get_dataset_id(dataset)
         metrics = self.validate(output, batch)
         for key, value in metrics.items():
             key = f"{dataset_id}/{key}"
@@ -290,7 +290,7 @@ def get_dataset(self, dataloader_idx: int) -> IRDataset | None:
             dataloaders = [dataloaders]
         return dataloaders[dataloader_idx].dataset
 
-    def get_dataset_id(self, dataloader_idx: int) -> str:
+    def get_dataset_id(self, dataset: IRDataset) -> str:
         """Gets the dataset id from the dataloader index for logging.
 
         .. _ir-datasets: https://ir-datasets.com/
@@ -300,9 +300,6 @@ def get_dataset_id(self, dataloader_idx: int) -> str:
         :return: path to run file, ir-datasets_ dataset id, or dataloader index
         :rtype: str
         """
-        dataset = self.get_dataset(dataloader_idx)
-        if dataset is None:
-            return str(dataloader_idx)
         if isinstance(dataset, RunDataset) and dataset.run_path is not None:
             dataset_id = dataset.run_path.name
         else:
@@ -420,7 +417,15 @@ def on_validation_end(self) -> None:
         df.columns.name = None
 
         # bring into correct order when skipping inference datasets
-        dataset_ids = [self.get_dataset_id(i) for i in range(df.shape[0])]
+        datamodule = getattr(self.trainer, "datamodule", None)
+        if datamodule is not None and hasattr(datamodule, "inference_datasets"):
+            inference_datasets = datamodule.inference_datasets
+            if len(inference_datasets) != df.shape[0]:
+                raise ValueError(
+                    "Number of inference datasets does not match number of dataloaders. "
+                    "Check if the dataloaders are correctly configured."
+                )
+            dataset_ids = [self.get_dataset_id(dataset) for dataset in inference_datasets]
         df = df.reindex(dataset_ids)
 
         trainer.print(df)

lightning_ir/data/__init__.py

@@ -6,12 +6,13 @@
 
 from .data import DocSample, IndexBatch, QuerySample, RankBatch, RankSample, SearchBatch, TrainBatch
 from .datamodule import LightningIRDataModule
-from .dataset import DocDataset, QueryDataset, RunDataset, TupleDataset
+from .dataset import DocDataset, IRDataset, QueryDataset, RunDataset, TupleDataset
 
 __all__ = [
     "DocDataset",
     "DocSample",
     "IndexBatch",
+    "IRDataset",
     "LightningIRDataModule",
     "QueryDataset",
     "QuerySample",