add docs_dataset_id to all indexes

webis-de · Dec 15, 2024 · 8b7c25c · 8b7c25c
1 parent 20d6489
commit 8b7c25c
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/lightning_ir/lightning_utils/callbacks.py b/lightning_ir/lightning_utils/callbacks.py
@@ -58,8 +58,9 @@ def get_index_dir(self, pl_module: BiEncoderModule, dataset: DocDataset) -> Path
                 index_dir = default_index_dir / "indexes"
             else:
                 raise ValueError("No index_dir provided and model_name_or_path is not a path")
-            index_dir = index_dir / dataset.docs_dataset_id
-        return Path(index_dir)
+        index_dir = Path(index_dir)
+        index_dir = index_dir / dataset.docs_dataset_id
+        return index_dir
 
 
 class _OverwriteMixin:
@@ -116,9 +117,8 @@ def get_indexer(self, trainer: Trainer, pl_module: BiEncoderModule, dataset_idx:
         dataloaders = trainer.test_dataloaders
         if dataloaders is None:
             raise ValueError("No test_dataloaders found")
-        dataset = dataloaders[dataset_idx].dataset
 
-        index_dir = self.get_index_dir(pl_module, dataset)
+        index_dir = self.get_save_path(trainer, pl_module, dataset_idx)
 
         indexer = self.index_config.indexer_class(index_dir, self.index_config, pl_module.config, self.verbose)
         return indexer