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

colpali-main/colpali_engine/trainer/retrieval_evaluator.py

@@ -53,14 +53,25 @@ class CustomEvaluator:
             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))
-            scores_batch = torch.cat(scores_batch, dim=1).cpu()
-            scores.append(scores_batch)
-        scores = torch.cat(scores, dim=0)
+            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
 
     def evaluate_biencoder(self, qs, ps) -> torch.Tensor: