Spaces:
Running
Running
File size: 3,120 Bytes
9ff79dc 7e4dcaf 9ff79dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import torch
from mteb.evaluation.evaluators import RetrievalEvaluator
class CustomEvaluator:
def __init__(self, is_multi_vector=False):
self.is_multi_vector = is_multi_vector
self.mteb_evaluator = RetrievalEvaluator()
def evaluate(self, qs, ps):
if self.is_multi_vector:
scores = self.evaluate_colbert(qs, ps)
else:
scores = self.evaluate_biencoder(qs, ps)
assert scores.shape[0] == len(qs)
arg_score = scores.argmax(dim=1)
# compare to arange
accuracy = (arg_score == torch.arange(scores.shape[0], device=scores.device)).sum().item() / scores.shape[0]
print(arg_score)
print(f"Top 1 Accuracy (verif): {accuracy}")
# cast to numpy
# scores = scores.cpu().numpy()
scores = scores.to(torch.float32).cpu().numpy()
return scores
def compute_metrics(self, relevant_docs, results, **kwargs):
# wrap mteb package
ndcg, _map, recall, precision, naucs = self.mteb_evaluator.evaluate(
relevant_docs,
results,
self.mteb_evaluator.k_values,
ignore_identical_ids=kwargs.get("ignore_identical_ids", True),
)
mrr = self.mteb_evaluator.evaluate_custom(relevant_docs, results, self.mteb_evaluator.k_values, "mrr")
scores = {
**{f"ndcg_at_{k.split('@')[1]}": v for (k, v) in ndcg.items()},
**{f"map_at_{k.split('@')[1]}": v for (k, v) in _map.items()},
**{f"recall_at_{k.split('@')[1]}": v for (k, v) in recall.items()},
**{f"precision_at_{k.split('@')[1]}": v for (k, v) in precision.items()},
**{f"mrr_at_{k.split('@')[1]}": v for (k, v) in mrr[0].items()},
**{f"naucs_at_{k.split('@')[1]}": v for (k, v) in naucs.items()},
}
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
def evaluate_biencoder(self, qs, ps) -> torch.Tensor:
qs = torch.stack(qs)
ps = torch.stack(ps)
scores = torch.einsum("bd,cd->bc", qs, ps)
return scores
|