refactor: use list_supported_models() (#501)

qdrant · Feb 21, 2024 · 0803fb5 · 0803fb5
1 parent fa91af2
commit 0803fb5
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 16 deletions.
diff --git a/qdrant_client/async_qdrant_fastembed.py b/qdrant_client/async_qdrant_fastembed.py
@@ -23,14 +23,14 @@
     from fastembed import TextEmbedding
 except ImportError:
     TextEmbedding = None
-SUPPORTED_EMBEDDING_MODELS: Dict[str, Tuple[int, models.Distance]] = {
-    "BAAI/bge-base-en": (768, models.Distance.COSINE),
-    "sentence-transformers/all-MiniLM-L6-v2": (384, models.Distance.COSINE),
-    "BAAI/bge-small-en": (384, models.Distance.COSINE),
-    "BAAI/bge-small-en-v1.5": (384, models.Distance.COSINE),
-    "BAAI/bge-base-en-v1.5": (768, models.Distance.COSINE),
-    "intfloat/multilingual-e5-large": (1024, models.Distance.COSINE),
-}
+SUPPORTED_EMBEDDING_MODELS: Dict[str, Tuple[int, models.Distance]] = (
+    {
+        model["model"]: (model["dim"], models.Distance.COSINE)
+        for model in TextEmbedding.list_supported_models()
+    }
+    if TextEmbedding
+    else {}
+)
 _DEPRECATED_MODELS = {"BAAI/bge-base-en", "BAAI/bge-small-en"}
 
 

diff --git a/qdrant_client/qdrant_fastembed.py b/qdrant_client/qdrant_fastembed.py
@@ -14,14 +14,14 @@
     TextEmbedding = None
 
 
-SUPPORTED_EMBEDDING_MODELS: Dict[str, Tuple[int, models.Distance]] = {
-    "BAAI/bge-base-en": (768, models.Distance.COSINE),
-    "sentence-transformers/all-MiniLM-L6-v2": (384, models.Distance.COSINE),
-    "BAAI/bge-small-en": (384, models.Distance.COSINE),
-    "BAAI/bge-small-en-v1.5": (384, models.Distance.COSINE),
-    "BAAI/bge-base-en-v1.5": (768, models.Distance.COSINE),
-    "intfloat/multilingual-e5-large": (1024, models.Distance.COSINE),
-}
+SUPPORTED_EMBEDDING_MODELS: Dict[str, Tuple[int, models.Distance]] = (
+    {
+        model["model"]: (model["dim"], models.Distance.COSINE)
+        for model in TextEmbedding.list_supported_models()
+    }
+    if TextEmbedding
+    else {}
+)
 
 _DEPRECATED_MODELS = {
     "BAAI/bge-base-en",