diff --git a/mteb/models/nomic_models.py b/mteb/models/nomic_models.py index 254d79dd5a..62fb40bdc9 100644 --- a/mteb/models/nomic_models.py +++ b/mteb/models/nomic_models.py @@ -4,6 +4,7 @@ from functools import partial from typing import Any +import numpy as np import torch import torch.nn.functional as F from sentence_transformers import SentenceTransformer @@ -43,7 +44,7 @@ def encode( # type: ignore prompt_type: PromptType | None = None, batch_size: int = 32, **kwargs: Any, - ): + ) -> np.ndarray: input_type = self.get_prompt_name(self.model_prompts, task_name, prompt_type) # default to search_document if input_type and prompt_name are not provided @@ -60,8 +61,9 @@ def encode( # type: ignore emb = torch.tensor(emb) emb = F.layer_norm(emb, normalized_shape=(emb.shape[1],)) emb = F.normalize(emb, p=2, dim=1) - if kwargs.get("convert_to_tensor", False): - emb = emb.cpu().detach().numpy() + + if isinstance(emb, torch.Tensor): + emb = emb.cpu().detach().float().numpy() return emb