feat: encoder types

semantic_router/encoders/__init__.py

@@ -1,7 +1,7 @@
 from typing import List, Optional
 
 from semantic_router.encoders.aurelio import AurelioSparseEncoder
-from semantic_router.encoders.base import BaseEncoder
+from semantic_router.encoders.base import DenseEncoder, SparseEncoder
 from semantic_router.encoders.bedrock import BedrockEncoder
 from semantic_router.encoders.bm25 import BM25Encoder
 from semantic_router.encoders.clip import CLIPEncoder
@@ -19,7 +19,8 @@
 
 __all__ = [
     "AurelioSparseEncoder",
-    "BaseEncoder",
+    "DenseEncoder",
+    "SparseEncoder",
     "AzureOpenAIEncoder",
     "CohereEncoder",
     "OpenAIEncoder",
@@ -39,7 +40,7 @@
 class AutoEncoder:
     type: EncoderType
     name: Optional[str]
-    model: BaseEncoder
+    model: DenseEncoder | SparseEncoder
 
     def __init__(self, type: str, name: Optional[str]):
         self.type = EncoderType(type)

semantic_router/encoders/aurelio.py

@@ -4,11 +4,11 @@
 
 from aurelio_sdk import AurelioClient, AsyncAurelioClient, EmbeddingResponse
 
-from semantic_router.encoders.base import BaseEncoder
+from semantic_router.encoders.base import SparseEncoder
 from semantic_router.schema import SparseEmbedding
 
 
-class AurelioSparseEncoder(BaseEncoder):
+class AurelioSparseEncoder(SparseEncoder):
     model: Optional[Any] = None
     idx_mapping: Optional[Dict[int, int]] = None
     client: AurelioClient = Field(default_factory=AurelioClient, exclude=True)

semantic_router/encoders/base.py

@@ -2,8 +2,10 @@
 
 from pydantic.v1 import BaseModel, Field, validator
 
+from semantic_router.schema import SparseEmbedding
 
-class BaseEncoder(BaseModel):
+
+class DenseEncoder(BaseModel):
     name: str
     score_threshold: Optional[float] = None
     type: str = Field(default="base")
@@ -20,3 +22,17 @@ def __call__(self, docs: List[Any]) -> List[List[float]]:
 
     def acall(self, docs: List[Any]) -> Coroutine[Any, Any, List[List[float]]]:
         raise NotImplementedError("Subclasses must implement this method")
+
+
+class SparseEncoder(BaseModel):
+    name: str
+    type: str = Field(default="base")
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def __call__(self, docs: List[str]) -> List[SparseEmbedding]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def acall(self, docs: List[str]) -> Coroutine[Any, Any, List[SparseEmbedding]]:
+        raise NotImplementedError("Subclasses must implement this method")

semantic_router/encoders/bedrock.py

@@ -1,7 +1,7 @@
 """
 This module provides the BedrockEncoder class for generating embeddings using Amazon's Bedrock Platform.
 
-The BedrockEncoder class is a subclass of BaseEncoder and utilizes the TextEmbeddingModel from the
+The BedrockEncoder class is a subclass of DenseEncoder and utilizes the TextEmbeddingModel from the
 Amazon's Bedrock Platform to generate embeddings for given documents. It requires an AWS Access Key ID
 and AWS Secret Access Key and supports customization of the pre-trained model, score threshold, and region.
 
@@ -21,12 +21,12 @@
 import os
 from time import sleep
 import tiktoken
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import DenseEncoder
 from semantic_router.utils.defaults import EncoderDefault
 from semantic_router.utils.logger import logger
 
 
-class BedrockEncoder(BaseEncoder):
+class BedrockEncoder(DenseEncoder):
     client: Any = None
     type: str = "bedrock"
     input_type: Optional[str] = "search_query"

semantic_router/encoders/bm25.py

@@ -1,10 +1,10 @@
 from typing import Any, Dict, List, Optional
 
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import SparseEncoder
 from semantic_router.utils.logger import logger
 
 
-class BM25Encoder(BaseEncoder):
+class BM25Encoder(SparseEncoder):
     model: Optional[Any] = None
     idx_mapping: Optional[Dict[int, int]] = None
     type: str = "sparse"

semantic_router/encoders/clip.py

@@ -3,10 +3,10 @@
 import numpy as np
 from pydantic.v1 import PrivateAttr
 from typing import Dict
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import DenseEncoder
 
 
-class CLIPEncoder(BaseEncoder):
+class CLIPEncoder(DenseEncoder):
     name: str = "openai/clip-vit-base-patch16"
     type: str = "huggingface"
     score_threshold: float = 0.2

semantic_router/encoders/cohere.py

@@ -3,11 +3,11 @@
 
 from pydantic.v1 import PrivateAttr
 
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import DenseEncoder
 from semantic_router.utils.defaults import EncoderDefault
 
 
-class CohereEncoder(BaseEncoder):
+class CohereEncoder(DenseEncoder):
     _client: Any = PrivateAttr()
     _embed_type: Any = PrivateAttr()
     type: str = "cohere"

semantic_router/encoders/fastembed.py

@@ -3,10 +3,10 @@
 import numpy as np
 from pydantic.v1 import PrivateAttr
 
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import DenseEncoder
 
 
-class FastEmbedEncoder(BaseEncoder):
+class FastEmbedEncoder(DenseEncoder):
     type: str = "fastembed"
     name: str = "BAAI/bge-small-en-v1.5"
     max_length: int = 512

semantic_router/encoders/google.py

@@ -1,7 +1,7 @@
 """
 This module provides the GoogleEncoder class for generating embeddings using Google's AI Platform.
 
-The GoogleEncoder class is a subclass of BaseEncoder and utilizes the TextEmbeddingModel from the
+The GoogleEncoder class is a subclass of DenseEncoder and utilizes the TextEmbeddingModel from the
 Google AI Platform to generate embeddings for given documents. It requires a Google Cloud project ID
 and supports customization of the pre-trained model, score threshold, location, and API endpoint.
 
@@ -19,11 +19,11 @@
 import os
 from typing import Any, List, Optional
 
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import DenseEncoder
 from semantic_router.utils.defaults import EncoderDefault
 
 
-class GoogleEncoder(BaseEncoder):
+class GoogleEncoder(DenseEncoder):
     """GoogleEncoder class for generating embeddings using Google's AI Platform.
 
     Attributes:

semantic_router/encoders/huggingface.py

@@ -1,7 +1,7 @@
 """
 This module provides the HFEndpointEncoder class to embeddings models using Huggingface's endpoint.
 
-The HFEndpointEncoder class is a subclass of BaseEncoder and utilizes a specified Huggingface 
+The HFEndpointEncoder class is a subclass of DenseEncoder and utilizes a specified Huggingface 
 endpoint to generate embeddings for given documents. It requires the URL of the Huggingface 
 API endpoint and an API key for authentication. The class supports customization of the score 
 threshold for filtering or processing the embeddings.
@@ -27,11 +27,11 @@
 
 from pydantic.v1 import PrivateAttr
 
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import DenseEncoder
 from semantic_router.utils.logger import logger
 
 
-class HuggingFaceEncoder(BaseEncoder):
+class HuggingFaceEncoder(DenseEncoder):
     name: str = "sentence-transformers/all-MiniLM-L6-v2"
     type: str = "huggingface"
     score_threshold: float = 0.5
@@ -140,7 +140,7 @@ def _max_pooling(self, model_output, attention_mask):
         return self._torch.max(token_embeddings, 1)[0]
 
 
-class HFEndpointEncoder(BaseEncoder):
+class HFEndpointEncoder(DenseEncoder):
     """
     A class to encode documents using a Hugging Face transformer model endpoint.
 

semantic_router/encoders/mistral.py

@@ -6,11 +6,11 @@
 
 from pydantic.v1 import PrivateAttr
 
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import DenseEncoder
 from semantic_router.utils.defaults import EncoderDefault
 
 
-class MistralEncoder(BaseEncoder):
+class MistralEncoder(DenseEncoder):
     """Class to encode text using MistralAI"""
 
     _client: Any = PrivateAttr()

semantic_router/encoders/openai.py

@@ -10,7 +10,7 @@
 from openai.types import CreateEmbeddingResponse
 import tiktoken
 
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import DenseEncoder
 from semantic_router.schema import EncoderInfo
 from semantic_router.utils.defaults import EncoderDefault
 from semantic_router.utils.logger import logger
@@ -35,7 +35,7 @@
 }
 
 
-class OpenAIEncoder(BaseEncoder):
+class OpenAIEncoder(DenseEncoder):
     client: Optional[openai.Client]
     async_client: Optional[openai.AsyncClient]
     dimensions: Union[int, NotGiven] = NotGiven()

semantic_router/encoders/tfidf.py

@@ -6,11 +6,11 @@
 from numpy import ndarray
 from numpy.linalg import norm
 
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import SparseEncoder
 from semantic_router.route import Route
 
 
-class TfidfEncoder(BaseEncoder):
+class TfidfEncoder(SparseEncoder):
     idf: ndarray = np.array([])
     word_index: Dict = {}
 

semantic_router/encoders/vit.py

@@ -2,10 +2,10 @@
 
 from pydantic.v1 import PrivateAttr
 
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import DenseEncoder
 
 
-class VitEncoder(BaseEncoder):
+class VitEncoder(DenseEncoder):
     name: str = "google/vit-base-patch16-224"
     type: str = "huggingface"
     score_threshold: float = 0.5

semantic_router/encoders/zure.py

@@ -8,12 +8,12 @@
 from openai import OpenAIError
 from openai.types import CreateEmbeddingResponse
 
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import DenseEncoder
 from semantic_router.utils.defaults import EncoderDefault
 from semantic_router.utils.logger import logger
 
 
-class AzureOpenAIEncoder(BaseEncoder):
+class AzureOpenAIEncoder(DenseEncoder):
     client: Optional[openai.AzureOpenAI] = None
     async_client: Optional[openai.AsyncAzureOpenAI] = None
     dimensions: Union[int, NotGiven] = NotGiven()

semantic_router/routers/base.py

@@ -10,7 +10,7 @@
 import yaml  # type: ignore
 from tqdm.auto import tqdm
 
-from semantic_router.encoders import AutoEncoder, BaseEncoder, OpenAIEncoder
+from semantic_router.encoders import AutoEncoder, DenseEncoder, OpenAIEncoder
 from semantic_router.index.base import BaseIndex
 from semantic_router.index.local import LocalIndex
 from semantic_router.index.pinecone import PineconeIndex
@@ -307,7 +307,7 @@ def get_hash(self) -> ConfigParameter:
 
 
 class BaseRouter(BaseModel):
-    encoder: BaseEncoder
+    encoder: DenseEncoder
     index: BaseIndex = Field(default_factory=BaseIndex)
     score_threshold: Optional[float] = Field(default=None)
     routes: List[Route] = []
@@ -322,7 +322,7 @@ class Config:
 
     def __init__(
         self,
-        encoder: Optional[BaseEncoder] = None,
+        encoder: Optional[DenseEncoder] = None,
         llm: Optional[BaseLLM] = None,
         routes: List[Route] = [],
         index: Optional[BaseIndex] = None,  # type: ignore

semantic_router/routers/hybrid.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from semantic_router.encoders import (
-    BaseEncoder,
+    DenseEncoder,
     BM25Encoder,
     TfidfEncoder,
 )
@@ -21,13 +21,13 @@ class HybridRouter(BaseRouter):
     """A hybrid layer that uses both dense and sparse embeddings to classify routes."""
 
     # there are a few additional attributes for hybrid
-    sparse_encoder: Optional[BaseEncoder] = Field(default=None)
+    sparse_encoder: Optional[DenseEncoder] = Field(default=None)
     alpha: float = 0.3
 
     def __init__(
         self,
-        encoder: BaseEncoder,
-        sparse_encoder: Optional[BaseEncoder] = None,
+        encoder: DenseEncoder,
+        sparse_encoder: Optional[DenseEncoder] = None,
         llm: Optional[BaseLLM] = None,
         routes: List[Route] = [],
         index: Optional[HybridLocalIndex] = None,
@@ -61,7 +61,7 @@ def __init__(
         if self.auto_sync:
             self._init_index_state()
 
-    def _set_sparse_encoder(self, sparse_encoder: Optional[BaseEncoder]):
+    def _set_sparse_encoder(self, sparse_encoder: Optional[DenseEncoder]):
         if sparse_encoder is None:
             logger.warning("No sparse_encoder provided. Using default BM25Encoder.")
             self.sparse_encoder = BM25Encoder()

semantic_router/routers/semantic.py

@@ -6,7 +6,7 @@
 import numpy as np
 from tqdm.auto import tqdm
 
-from semantic_router.encoders import AutoEncoder, BaseEncoder, OpenAIEncoder
+from semantic_router.encoders import AutoEncoder, DenseEncoder, OpenAIEncoder
 from semantic_router.index.base import BaseIndex
 from semantic_router.index.local import LocalIndex
 from semantic_router.index.pinecone import PineconeIndex
@@ -55,7 +55,7 @@ def is_valid(layer_config: str) -> bool:
 class SemanticRouter(BaseRouter):
     def __init__(
         self,
-        encoder: Optional[BaseEncoder] = None,
+        encoder: Optional[DenseEncoder] = None,
         llm: Optional[BaseLLM] = None,
         routes: Optional[List[Route]] = None,
         index: Optional[BaseIndex] = None,  # type: ignore

tests/integration/encoders/test_openai_integration.py

@@ -1,7 +1,7 @@
 import os
 import pytest
 from openai import OpenAIError
-from semantic_router.encoders.base import BaseEncoder
+from semantic_router.encoders.base import DenseEncoder
 from semantic_router.encoders.openai import OpenAIEncoder
 
 with open("tests/integration/57640.4032.txt", "r") as fp:
@@ -11,7 +11,7 @@
 @pytest.fixture
 def openai_encoder():
     if os.environ.get("OPENAI_API_KEY") is None:
-        return BaseEncoder()
+        return DenseEncoder()
     else:
         return OpenAIEncoder()
 

tests/unit/encoders/test_base.py

@@ -1,12 +1,12 @@
 import pytest
 
-from semantic_router.encoders import BaseEncoder
+from semantic_router.encoders import DenseEncoder
 
 
-class TestBaseEncoder:
+class TestDenseEncoder:
     @pytest.fixture
     def base_encoder(self):
-        return BaseEncoder(name="TestEncoder", score_threshold=0.5)
+        return DenseEncoder(name="TestEncoder", score_threshold=0.5)
 
     def test_base_encoder_initialization(self, base_encoder):
         assert base_encoder.name == "TestEncoder", "Initialization of name failed"