Update colpali-main/colpali_engine/trainer/retrieval_evaluator.py

colpali-main/colpali_engine/trainer/retrieval_evaluator.py

@@ -47,32 +47,25 @@ class CustomEvaluator:
         return scores
 
     def evaluate_colbert(self, qs, ps, batch_size=128) -> torch.Tensor:
-        scores = []
-        for i in range(0, len(qs), batch_size):
-            scores_batch = []
-            qs_batch = torch.nn.utils.rnn.pad_sequence(qs[i : i + batch_size], batch_first=True, padding_value=0).to(
-                "cpu"
-            )
-            print(f"qs_batch shape: {qs_batch.shape}")  # Impression pour le débogage
-            for j in range(0, len(ps), batch_size):
-                ps_batch = torch.nn.utils.rnn.pad_sequence(
-                    ps[j : j + batch_size], batch_first=True, padding_value=0
-                ).to("cpu")
-                print(f"ps_batch shape: {ps_batch.shape}")  # Impression pour le débogage
-                scores_batch.append(torch.einsum("bnd,csd->bcns", qs_batch, ps_batch).max(dim=3)[0].sum(dim=2))
-            if scores_batch:  # Vérification si scores_batch n'est pas vide
-                print(f"scores_batch shapes before concat: {[s.shape for s in scores_batch]}")  # Impression pour le débogage
-                scores_batch = torch.cat(scores_batch, dim=1).cpu()
-                scores.append(scores_batch)
-            else:
-                print("scores_batch is empty!")  # Impression pour le débogage
-        if scores:  # Vérification si scores n'est pas vide
-            print(f"scores shapes before final concat: {[s.shape for s in scores]}")  # Impression pour le débogage
-            scores = torch.cat(scores, dim=0)
-        else:
-            print("scores is empty!")  # Impression pour le débogage
-            scores = torch.tensor([])  # Retourne un tensor vide si scores est vide
-        return scores
+    scores = []
+    for i in range(0, len(qs), batch_size):
+        scores_batch = []
+        qs_batch = torch.nn.utils.rnn.pad_sequence(qs[i : i + batch_size], batch_first=True, padding_value=0).to(
+            "cpu"
+        )
+        for j in range(0, len(ps), batch_size):
+            ps_batch = torch.nn.utils.rnn.pad_sequence(
+                ps[j : j + batch_size], batch_first=True, padding_value=0
+            ).to("cpu")
+            scores_batch.append(torch.einsum("bnd,csd->bcns", qs_batch, ps_batch).max(dim=3)[0].sum(dim=2))
+        if scores_batch:  # Vérification si scores_batch n'est pas vide
+            scores_batch = torch.cat(scores_batch, dim=1).cpu()
+            scores.append(scores_batch)
+    if scores:  # Vérification si scores n'est pas vide
+        scores = torch.cat(scores, dim=0)
+    else:
+        scores = torch.tensor([])  # Retourne un tensor vide si scores est vide
+    return scores
 
     def evaluate_biencoder(self, qs, ps) -> torch.Tensor: