From 7862efdbb09da27aa05f76e231a9a23a70abf583 Mon Sep 17 00:00:00 2001
From: Ismail Ashraq <issey1455@gmail.com>
Date: Sun, 7 Jan 2024 18:35:39 +0500
Subject: [PATCH] Move score_threshold to encoders
---
semantic_router/encoders/base.py | 1 +
semantic_router/encoders/bm25.py | 6 ++++--
semantic_router/encoders/cohere.py | 3 ++-
semantic_router/encoders/fastembed.py | 6 ++++--
semantic_router/encoders/openai.py | 3 ++-
semantic_router/layer.py | 20 +++++---------------
tests/unit/encoders/test_base.py | 3 ++-
tests/unit/test_layer.py | 7 +++++--
8 files changed, 25 insertions(+), 24 deletions(-)
diff --git a/semantic_router/encoders/base.py b/semantic_router/encoders/base.py
index bd952403..f5968578 100644
--- a/semantic_router/encoders/base.py
+++ b/semantic_router/encoders/base.py
@@ -3,6 +3,7 @@ from pydantic import BaseModel, Field
class BaseEncoder(BaseModel):
name: str
+ score_threshold: float
type: str = Field(default="base")
class Config:
diff --git a/semantic_router/encoders/bm25.py b/semantic_router/encoders/bm25.py
index e597b4fe..68150cb7 100644
--- a/semantic_router/encoders/bm25.py
+++ b/semantic_router/encoders/bm25.py
@@ -1,6 +1,7 @@
from typing import Any
from semantic_router.encoders import BaseEncoder
+from semantic_router.utils.logger import logger
class BM25Encoder(BaseEncoder):
@@ -8,8 +9,8 @@ class BM25Encoder(BaseEncoder):
idx_mapping: dict[int, int] | None = None
type: str = "sparse"
- def __init__(self, name: str = "bm25"):
- super().__init__(name=name)
+ def __init__(self, name: str = "bm25", score_threshold: float = 0.82):
+ super().__init__(name=name, score_threshold=score_threshold)
try:
from pinecone_text.sparse import BM25Encoder as encoder
except ImportError:
@@ -17,6 +18,7 @@ class BM25Encoder(BaseEncoder):
"Please install pinecone-text to use BM25Encoder. "
"You can install it with: `pip install semantic-router[hybrid]`"
)
+ logger.info("Downloading and initializing BM25 model parameters.")
self.model = encoder.default()
params = self.model.get_params()
diff --git a/semantic_router/encoders/cohere.py b/semantic_router/encoders/cohere.py
index f7aef0e6..2f80aaaf 100644
--- a/semantic_router/encoders/cohere.py
+++ b/semantic_router/encoders/cohere.py
@@ -13,10 +13,11 @@ class CohereEncoder(BaseEncoder):
self,
name: str | None = None,
cohere_api_key: str | None = None,
+ score_threshold: float = 0.3,
):
if name is None:
name = os.getenv("COHERE_MODEL_NAME", "embed-english-v3.0")
- super().__init__(name=name)
+ super().__init__(name=name, score_threshold=score_threshold)
cohere_api_key = cohere_api_key or os.getenv("COHERE_API_KEY")
if cohere_api_key is None:
raise ValueError("Cohere API key cannot be 'None'.")
diff --git a/semantic_router/encoders/fastembed.py b/semantic_router/encoders/fastembed.py
index fb845ce7..413e3a6a 100644
--- a/semantic_router/encoders/fastembed.py
+++ b/semantic_router/encoders/fastembed.py
@@ -14,8 +14,10 @@ class FastEmbedEncoder(BaseEncoder):
threads: Optional[int] = None
_client: Any = PrivateAttr()
- def __init__(self, **data):
- super().__init__(**data)
+ def __init__(
+ self, score_threshold: float = 0.5, **data
+ ): # TODO default score_threshold not thoroughly tested, should optimize
+ super().__init__(score_threshold=score_threshold, **data)
self._client = self._initialize_client()
def _initialize_client(self):
diff --git a/semantic_router/encoders/openai.py b/semantic_router/encoders/openai.py
index f9348a12..4ec87638 100644
--- a/semantic_router/encoders/openai.py
+++ b/semantic_router/encoders/openai.py
@@ -17,10 +17,11 @@ class OpenAIEncoder(BaseEncoder):
self,
name: str | None = None,
openai_api_key: str | None = None,
+ score_threshold: float = 0.82,
):
if name is None:
name = os.getenv("OPENAI_MODEL_NAME", "text-embedding-ada-002")
- super().__init__(name=name)
+ super().__init__(name=name, score_threshold=score_threshold)
api_key = openai_api_key or os.getenv("OPENAI_API_KEY")
if api_key is None:
raise ValueError("OpenAI API key cannot be 'None'.")
diff --git a/semantic_router/layer.py b/semantic_router/layer.py
index 6d85508c..a71411db 100644
--- a/semantic_router/layer.py
+++ b/semantic_router/layer.py
@@ -7,8 +7,6 @@ import yaml
from semantic_router.encoders import (
BaseEncoder,
CohereEncoder,
- OpenAIEncoder,
- FastEmbedEncoder,
)
from semantic_router.linear import similarity_matrix, top_scores
from semantic_router.route import Route
@@ -153,27 +151,19 @@ class LayerConfig:
class RouteLayer:
index: np.ndarray | None = None
categories: np.ndarray | None = None
- score_threshold: float = 0.82
+ score_threshold: float
def __init__(
- self, encoder: BaseEncoder | None = None, routes: list[Route] | None = None
+ self,
+ encoder: BaseEncoder | None = None,
+ routes: list[Route] | None = None,
):
logger.info("Initializing RouteLayer")
self.index = None
self.categories = None
self.encoder = encoder if encoder is not None else CohereEncoder()
self.routes: list[Route] = routes if routes is not None else []
- # decide on default threshold based on encoder
- # TODO move defaults to the encoder objects and extract from there
- if isinstance(encoder, OpenAIEncoder):
- self.score_threshold = 0.82
- elif isinstance(encoder, CohereEncoder):
- self.score_threshold = 0.3
- elif isinstance(encoder, FastEmbedEncoder):
- # TODO default not thoroughly tested, should optimize
- self.score_threshold = 0.5
- else:
- self.score_threshold = 0.82
+ self.score_threshold = self.encoder.score_threshold
# if routes list has been passed, we initialize index now
if len(self.routes) > 0:
# initialize index now
diff --git a/tests/unit/encoders/test_base.py b/tests/unit/encoders/test_base.py
index d2c39645..4d4b87ae 100644
--- a/tests/unit/encoders/test_base.py
+++ b/tests/unit/encoders/test_base.py
@@ -6,10 +6,11 @@ from semantic_router.encoders import BaseEncoder
class TestBaseEncoder:
@pytest.fixture
def base_encoder(self):
- return BaseEncoder(name="TestEncoder")
+ return BaseEncoder(name="TestEncoder", score_threshold=0.5)
def test_base_encoder_initialization(self, base_encoder):
assert base_encoder.name == "TestEncoder", "Initialization of name failed"
+ assert base_encoder.score_threshold == 0.5
def test_base_encoder_call_method_not_implemented(self, base_encoder):
with pytest.raises(NotImplementedError):
diff --git a/tests/unit/test_layer.py b/tests/unit/test_layer.py
index 32754997..6652579f 100644
--- a/tests/unit/test_layer.py
+++ b/tests/unit/test_layer.py
@@ -69,7 +69,7 @@ routes:
@pytest.fixture
def base_encoder():
- return BaseEncoder(name="test-encoder")
+ return BaseEncoder(name="test-encoder", score_threshold=0.5)
@pytest.fixture
@@ -95,6 +95,7 @@ def routes():
class TestRouteLayer:
def test_initialization(self, openai_encoder, routes):
route_layer = RouteLayer(encoder=openai_encoder, routes=routes)
+ assert openai_encoder.score_threshold == 0.82
assert route_layer.score_threshold == 0.82
assert len(route_layer.index) if route_layer.index is not None else 0 == 5
assert (
@@ -105,9 +106,11 @@ class TestRouteLayer:
def test_initialization_different_encoders(self, cohere_encoder, openai_encoder):
route_layer_cohere = RouteLayer(encoder=cohere_encoder)
+ assert cohere_encoder.score_threshold == 0.3
assert route_layer_cohere.score_threshold == 0.3
route_layer_openai = RouteLayer(encoder=openai_encoder)
+ assert openai_encoder.score_threshold == 0.82
assert route_layer_openai.score_threshold == 0.82
def test_add_route(self, openai_encoder):
@@ -173,7 +176,7 @@ class TestRouteLayer:
def test_failover_score_threshold(self, base_encoder):
route_layer = RouteLayer(encoder=base_encoder)
- assert route_layer.score_threshold == 0.82
+ assert route_layer.score_threshold == 0.5
def test_json(self, openai_encoder, routes):
with tempfile.NamedTemporaryFile(suffix=".yaml", delete=False) as temp:
--
GitLab