feat: Add fastembed encoder

pyproject.toml

@@ -13,7 +13,7 @@ readme = "README.md"
 packages = [{include = "semantic_router"}]
 
 [tool.poetry.dependencies]
-python = "^3.9"
+python = ">=3.9,<3.12"
 pydantic = "^1.8.2"
 openai = "^1.3.9"
 cohere = "^4.32"
@@ -22,6 +22,7 @@ pinecone-text = "^0.7.0"
 colorlog = "^6.8.0"
 pyyaml = "^6.0.1"
 pytest-asyncio = "^0.23.2"
+fastembed = "^0.1.3"
 
 
 [tool.poetry.group.dev.dependencies]

semantic_router/encoders/__init__.py

@@ -2,5 +2,12 @@
 from semantic_router.encoders.bm25 import BM25Encoder
 from semantic_router.encoders.cohere import CohereEncoder
 from semantic_router.encoders.openai import OpenAIEncoder
+from semantic_router.encoders.fastembed import FastEmbedEncoder
 
-__all__ = ["BaseEncoder", "CohereEncoder", "OpenAIEncoder", "BM25Encoder"]
+__all__ = [
+    "BaseEncoder",
+    "CohereEncoder",
+    "OpenAIEncoder",
+    "BM25Encoder",
+    "FastEmbedEncoder",
+]

semantic_router/encoders/fastembed.py

@@ -0,0 +1,45 @@
+from typing import Any, List, Optional
+import numpy as np
+from pydantic import BaseModel, PrivateAttr
+
+
+class FastEmbedEncoder(BaseModel):
+    type: str = "fastembed"
+    model_name: str = "BAAI/bge-small-en-v1.5"
+    max_length: int = 512
+    cache_dir: Optional[str] = None
+    threads: Optional[int] = None
+    _client: Any = PrivateAttr()
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self._client = self._initialize_client()
+
+    def _initialize_client(self):
+        try:
+            from fastembed.embedding import FlagEmbedding as Embedding
+        except ImportError:
+            raise ImportError(
+                "Please install fastembed to use FastEmbedEncoder"
+                "You can install it with: `pip install fastembed`"
+            )
+
+        embedding_args = {
+            "model_name": self.model_name,
+            "max_length": self.max_length,
+            "cache_dir": self.cache_dir,
+            "threads": self.threads,
+        }
+
+        embedding_args = {k: v for k, v in embedding_args.items() if v is not None}
+
+        embedding = Embedding(**embedding_args)
+        return embedding
+
+    def __call__(self, docs: list[str]) -> list[list[float]]:
+        try:
+            embeds: List[np.ndarray] = list(self._client.embed(docs))
+            embeddings: List[List[float]] = [e.tolist() for e in embeds]
+            return embeddings
+        except Exception as e:
+            raise ValueError(f"FastEmbed embed failed. Error: {e}")

tests/unit/encoders/test_fastembed.py

@@ -0,0 +1,10 @@
+from semantic_router.encoders import FastEmbedEncoder
+
+
+class TestFastEmbedEncoder:
+    def test_fastembed_encoder(self):
+        encode = FastEmbedEncoder()
+        test_docs = ["This is a test", "This is another test"]
+
+        embeddings = encode(test_docs)
+        assert isinstance(embeddings, list)