Merge pull request #61 from fschlatt/main

Fix CI + Minor pretty printing fix

Author: fschlatt

lightning_ir/base/module.py

@@ -14,7 +14,7 @@
 from lightning.pytorch.trainer.states import RunningStage
 from transformers import BatchEncoding
 
-from ..data import LightningIRDataModule, RankBatch, RunDataset, SearchBatch, TrainBatch
+from ..data import RankBatch, RunDataset, SearchBatch, TrainBatch
 from ..data.dataset import IRDataset
 from ..loss.loss import InBatchLossFunction, LossFunction
 from .config import LightningIRConfig
@@ -414,10 +414,9 @@ def on_validation_end(self) -> None:
         df = df.pivot(index="dataset", columns="metric", values="value")
         df.columns.name = None
 
-        datamodule: LightningIRDataModule | None = getattr(trainer, "datamodule", None)
-        if datamodule is not None and datamodule.inference_datasets is not None:
-            dataset_ids = [dataset.dataset_id for dataset in datamodule.inference_datasets]
-            df = df.reindex(dataset_ids)
+        # bring into correct order when skipping inference datasets
+        dataset_ids = [self.get_dataset_id(i) for i in range(df.shape[0])]
+        df = df.reindex(dataset_ids)
 
         trainer.print(df)
 

setup.cfg

@@ -42,6 +42,7 @@ test =
     sentence-transformers
     faiss-cpu
     pyseismic-lsr
+    pylate
 dev =
     black
     flake8

tests/test_models/test_col.py

@@ -5,7 +5,6 @@
 
 from colbert.modeling.checkpoint import Checkpoint  # noqa: E402
 from colbert.modeling.colbert import ColBERTConfig, colbert_score  # noqa: E402
-
 from pylate import models, rank
 
 from lightning_ir import BiEncoderModule  # noqa: E402
@@ -55,16 +54,16 @@ def test_same_as_modern_colbert():
     doc_embedding = output.doc_embeddings
 
     orig_model = models.ColBERT(model_name_or_path=model_name)
-    orig_query = orig_model.encode(
-        [query],
-        is_query=True,
-    )
-    orig_docs = orig_model.encode(
-        [documents],
-        is_query=False,
+    orig_query = orig_model.encode([query], is_query=True)
+    orig_docs = orig_model.encode([documents], is_query=False)
+    orig_scores = rank.rerank(
+        queries_embeddings=orig_query, documents_embeddings=orig_docs, documents_ids=[list(range(len(documents)))]
     )
-    orig_scores = rank.rerank(queries_embeddings=orig_query, documents_embeddings=orig_docs, documents_ids=[list(range(len(documents)))])
 
     assert torch.allclose(query_embedding.embeddings, torch.tensor(orig_query[0]), atol=1e-6)
-    assert torch.allclose(doc_embedding.embeddings[doc_embedding.scoring_mask], torch.cat([torch.from_numpy(d) for doc in orig_docs for d in doc]), atol=1e-6)
-    assert torch.allclose(output.scores, torch.tensor([d["score"] for q in orig_scores for d in q]), atol=1e-6)
+    assert torch.allclose(
+        doc_embedding.embeddings[doc_embedding.scoring_mask],
+        torch.cat([torch.from_numpy(d) for doc in orig_docs for d in doc]),
+        atol=1e-6,
+    )
+    assert torch.allclose(output.scores, torch.tensor([d["score"] for q in orig_scores for d in q]), atol=1e-6)