From 885a738751908d9e917e17bec9e2cd2f6dcd8bf3 Mon Sep 17 00:00:00 2001 From: NirantK Date: Thu, 10 Aug 2023 10:43:50 +0530 Subject: [PATCH] * feat(embedding.py): add support for custom model name, onnx providers, and max length in DefaultEmbedding constructor --- fastembed/embedding.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/fastembed/embedding.py b/fastembed/embedding.py index 313e61b..ba06090 100644 --- a/fastembed/embedding.py +++ b/fastembed/embedding.py @@ -37,7 +37,10 @@ class ONNXProviders: Metal = "CoreMLExecutionProvider" class DefaultEmbedding(Embedding): - def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"): + def __init__(self, + model_name: str = "BAAI/bge-base-en", + onnx_providers: List[str] = [ONNXProviders.Metal], + max_length: int = 512): self.cache_dir = Path(tempfile.gettempdir()) / "fastembed" filepath = self.download_file_from_gcs( "https://storage.googleapis.com/qdrant-fastembed/fast-all-MiniLM-L6-v2.tar.gz",