From 34f6037ce0093c74ff7d3c7ce1355bf92350936d Mon Sep 17 00:00:00 2001 From: Anush Date: Thu, 22 Feb 2024 00:32:56 +0530 Subject: [PATCH] refactor: use list_supported_models() (#501) --- qdrant_client/async_qdrant_fastembed.py | 16 ++++++++-------- qdrant_client/qdrant_fastembed.py | 16 ++++++++-------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/qdrant_client/async_qdrant_fastembed.py b/qdrant_client/async_qdrant_fastembed.py index 4d103e01..167358a8 100644 --- 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 index e300be8e..3b8d7cb9 100644 --- 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",