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

colpali-main/colpali_engine/trainer/retrieval_evaluator.py

@@ -47,25 +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"
-        )
-        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
+        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: