From 7be2ca6bf82b0522238c9babfeeb71ec99964661 Mon Sep 17 00:00:00 2001
From: Roman Solomatin <36135455+Samoed@users.noreply.github.com>
Date: Fri, 21 Feb 2025 11:52:51 +0300
Subject: [PATCH] fix voyage

---
 mteb/models/voyage_models.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mteb/models/voyage_models.py b/mteb/models/voyage_models.py
index 9c5860ba66..d576288c3a 100644
--- a/mteb/models/voyage_models.py
+++ b/mteb/models/voyage_models.py
@@ -98,7 +98,8 @@ def encode(
         prompt_type: PromptType | None = None,
         **kwargs: Any,
     ) -> np.ndarray:
-        input_type = self.model_prompts.get(prompt_type.value, "document")
+        prompt_name = self.get_prompt_name(self.model_prompts, task_name, prompt_type)
+        input_type = prompt_name if prompt_name is not None else "document"
         return self._batched_encode(sentences, batch_size, input_type)
 
     def _batched_encode(