Spaces:
Sleeping
Sleeping
HUANG-Stephanie
commited on
Commit
•
7e4dcaf
1
Parent(s):
2b7ded4
Update colpali-main/colpali_engine/trainer/retrieval_evaluator.py
Browse files
colpali-main/colpali_engine/trainer/retrieval_evaluator.py
CHANGED
@@ -47,25 +47,25 @@ class CustomEvaluator:
|
|
47 |
return scores
|
48 |
|
49 |
def evaluate_colbert(self, qs, ps, batch_size=128) -> torch.Tensor:
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
|
56 |
-
|
57 |
-
|
58 |
-
|
59 |
-
|
60 |
-
|
61 |
-
|
62 |
-
|
63 |
-
|
64 |
-
|
65 |
-
|
66 |
-
|
67 |
-
|
68 |
-
|
69 |
|
70 |
def evaluate_biencoder(self, qs, ps) -> torch.Tensor:
|
71 |
|
|
|
47 |
return scores
|
48 |
|
49 |
def evaluate_colbert(self, qs, ps, batch_size=128) -> torch.Tensor:
|
50 |
+
scores = []
|
51 |
+
for i in range(0, len(qs), batch_size):
|
52 |
+
scores_batch = []
|
53 |
+
qs_batch = torch.nn.utils.rnn.pad_sequence(qs[i : i + batch_size], batch_first=True, padding_value=0).to(
|
54 |
+
"cpu"
|
55 |
+
)
|
56 |
+
for j in range(0, len(ps), batch_size):
|
57 |
+
ps_batch = torch.nn.utils.rnn.pad_sequence(
|
58 |
+
ps[j : j + batch_size], batch_first=True, padding_value=0
|
59 |
+
).to("cpu")
|
60 |
+
scores_batch.append(torch.einsum("bnd,csd->bcns", qs_batch, ps_batch).max(dim=3)[0].sum(dim=2))
|
61 |
+
if scores_batch: # Vérification si scores_batch n'est pas vide
|
62 |
+
scores_batch = torch.cat(scores_batch, dim=1).cpu()
|
63 |
+
scores.append(scores_batch)
|
64 |
+
if scores: # Vérification si scores n'est pas vide
|
65 |
+
scores = torch.cat(scores, dim=0)
|
66 |
+
else:
|
67 |
+
scores = torch.tensor([]) # Retourne un tensor vide si scores est vide
|
68 |
+
return scores
|
69 |
|
70 |
def evaluate_biencoder(self, qs, ps) -> torch.Tensor:
|
71 |
|