add new OpenAI model names (#7297)

deepset-ai · Mar 4, 2024 · d03d5b5 · d03d5b5
1 parent 5b59fcb
commit d03d5b5
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py
@@ -1895,9 +1895,14 @@ def _preprocess_documents(self, docs: List[Document]) -> List[Document]:
     @staticmethod
     def _infer_model_format(model_name_or_path: str, use_auth_token: Optional[Union[str, bool]]) -> str:
         # pylint: disable=too-many-return-statements
-        valid_openai_model_name = model_name_or_path in ["ada", "babbage", "davinci", "curie"] or any(
-            m in model_name_or_path for m in ["-ada-", "-babbage-", "-davinci-", "-curie-"]
-        )
+        valid_openai_model_name = model_name_or_path in [
+            "ada",
+            "babbage",
+            "davinci",
+            "curie",
+            "text-embedding-3-small",
+            "text-embedding-3-large",
+        ] or any(m in model_name_or_path for m in ["-ada-", "-babbage-", "-davinci-", "-curie-"])
         if valid_openai_model_name:
             return "openai"
         if model_name_or_path in COHERE_EMBEDDING_MODELS: