diff --git a/colpali_engine/trainer/retrieval_evaluator.py b/colpali_engine/trainer/retrieval_evaluator.py index 9681b29d..16a386cd 100644 --- a/colpali_engine/trainer/retrieval_evaluator.py +++ b/colpali_engine/trainer/retrieval_evaluator.py @@ -1,11 +1,12 @@ import torch from mteb.evaluation.evaluators import RetrievalEvaluator - +from colpali_engine.utils.torch_utils import get_torch_device class CustomEvaluator: def __init__(self, is_multi_vector=False): self.is_multi_vector = is_multi_vector self.mteb_evaluator = RetrievalEvaluator() + self.device = get_torch_device() def evaluate(self, qs, ps): if self.is_multi_vector: @@ -51,12 +52,12 @@ def evaluate_colbert(self, qs, ps, batch_size=128) -> torch.Tensor: 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( - "cuda" + self.device ) 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("cuda") + ).to(self.device) 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)