Skip to content
Snippets Groups Projects
Commit d5f4703d authored by jamescalam's avatar jamescalam
Browse files

feat: encoder types

parent 87d6ae4b
Branches
Tags
No related merge requests found
Showing
with 66 additions and 49 deletions
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 @@ from semantic_router.schema import EncoderType
__all__ = [
"AurelioSparseEncoder",
"BaseEncoder",
"DenseEncoder",
"SparseEncoder",
"AzureOpenAIEncoder",
"CohereEncoder",
"OpenAIEncoder",
......@@ -39,7 +40,7 @@ __all__ = [
class AutoEncoder:
type: EncoderType
name: Optional[str]
model: BaseEncoder
model: DenseEncoder | SparseEncoder
def __init__(self, type: str, name: Optional[str]):
self.type = EncoderType(type)
......
......@@ -4,11 +4,11 @@ from pydantic.v1 import Field
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)
......
......@@ -2,8 +2,10 @@ from typing import Any, Coroutine, List, Optional
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 @@ class BaseEncoder(BaseModel):
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")
\ No newline at end of file
"""
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 @@ from typing import List, Optional, Any
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"
......
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"
......
......@@ -3,10 +3,10 @@ from typing import Any, List, Optional
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
......
......@@ -3,11 +3,11 @@ from typing import Any, List, Optional
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"
......
......@@ -3,10 +3,10 @@ from typing import Any, List, Optional
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
......
"""
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 @@ Classes:
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:
......
"""
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 typing import Any, List, Optional, Dict
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 @@ class HuggingFaceEncoder(BaseEncoder):
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.
......
......@@ -6,11 +6,11 @@ from typing import Any, List, Optional
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()
......
......@@ -10,7 +10,7 @@ from openai._types import NotGiven
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 @@ model_configs = {
}
class OpenAIEncoder(BaseEncoder):
class OpenAIEncoder(DenseEncoder):
client: Optional[openai.Client]
async_client: Optional[openai.AsyncClient]
dimensions: Union[int, NotGiven] = NotGiven()
......
......@@ -6,11 +6,11 @@ import numpy as np
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 = {}
......
......@@ -2,10 +2,10 @@ from typing import Any, List, Optional, Dict
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
......
......@@ -8,12 +8,12 @@ from openai._types import NotGiven
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()
......
......@@ -10,7 +10,7 @@ import numpy as np
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 @@ class RouterConfig:
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 BaseRouter(BaseModel):
def __init__(
self,
encoder: Optional[BaseEncoder] = None,
encoder: Optional[DenseEncoder] = None,
llm: Optional[BaseLLM] = None,
routes: List[Route] = [],
index: Optional[BaseIndex] = None, # type: ignore
......
......@@ -5,7 +5,7 @@ from pydantic.v1 import Field
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 @@ class HybridRouter(BaseRouter):
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()
......
......@@ -6,7 +6,7 @@ from pydantic.v1 import validator, Field
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
......
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 @@ with open("tests/integration/57640.4032.txt", "r") as fp:
@pytest.fixture
def openai_encoder():
if os.environ.get("OPENAI_API_KEY") is None:
return BaseEncoder()
return DenseEncoder()
else:
return OpenAIEncoder()
......
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"
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment