From 7862efdbb09da27aa05f76e231a9a23a70abf583 Mon Sep 17 00:00:00 2001 From: Ismail Ashraq Date: Sun, 7 Jan 2024 18:35:39 +0500 Subject: [PATCH 1/7] 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 @@ 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 @@ def __init__(self, name: str = "bm25"): "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 @@ def __init__( 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 @@ def __init__( 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 @@ 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 @@ def remove(self, name: str): 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 @@ 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 @@ def layer_yaml(): @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 @@ def test_initialization(self, openai_encoder, routes): 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 @@ def test_pass_threshold(self, openai_encoder): 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: From d551201404a2ed2c67e989bb55487d2f5a2f976e Mon Sep 17 00:00:00 2001 From: Ismail Ashraq Date: Sun, 7 Jan 2024 19:15:03 +0500 Subject: [PATCH 2/7] Update hybrid layer to use score_threshold from encoders --- semantic_router/encoders/bm25.py | 6 +++++- semantic_router/hybrid_layer.py | 12 ++---------- tests/unit/test_hybrid_layer.py | 6 ++++-- 3 files changed, 11 insertions(+), 13 deletions(-) diff --git a/semantic_router/encoders/bm25.py b/semantic_router/encoders/bm25.py index 68150cb7..69ca58ec 100644 --- a/semantic_router/encoders/bm25.py +++ b/semantic_router/encoders/bm25.py @@ -19,7 +19,11 @@ def __init__(self, name: str = "bm25", score_threshold: float = 0.82): "You can install it with: `pip install semantic-router[hybrid]`" ) logger.info("Downloading and initializing BM25 model parameters.") - self.model = encoder.default() + # self.model = encoder.default() + self.model = encoder() + self.model.fit( + corpus=["test test", "this is another message", "hello how are you"] + ) params = self.model.get_params() doc_freq = params["doc_freq"] diff --git a/semantic_router/hybrid_layer.py b/semantic_router/hybrid_layer.py index fc63cfa6..cd9f7ccb 100644 --- a/semantic_router/hybrid_layer.py +++ b/semantic_router/hybrid_layer.py @@ -4,8 +4,6 @@ from semantic_router.encoders import ( BaseEncoder, BM25Encoder, - CohereEncoder, - OpenAIEncoder, ) from semantic_router.route import Route from semantic_router.utils.logger import logger @@ -15,21 +13,15 @@ class HybridRouteLayer: index = None sparse_index = None categories = None - score_threshold = 0.82 + score_threshold: float def __init__( self, encoder: BaseEncoder, routes: list[Route] = [], alpha: float = 0.3 ): self.encoder = encoder + self.score_threshold = self.encoder.score_threshold self.sparse_encoder = BM25Encoder() self.alpha = alpha - # decide on default threshold based on encoder - if isinstance(encoder, OpenAIEncoder): - self.score_threshold = 0.82 - elif isinstance(encoder, CohereEncoder): - self.score_threshold = 0.3 - else: - self.score_threshold = 0.82 # if routes list has been passed, we initialize index now if routes: # initialize index now diff --git a/tests/unit/test_hybrid_layer.py b/tests/unit/test_hybrid_layer.py index f87cb1d2..6896c4de 100644 --- a/tests/unit/test_hybrid_layer.py +++ b/tests/unit/test_hybrid_layer.py @@ -19,7 +19,7 @@ def mock_encoder_call(utterances): @pytest.fixture def base_encoder(): - return BaseEncoder(name="test-encoder") + return BaseEncoder(name="test-encoder", score_threshold=0.5) @pytest.fixture @@ -46,6 +46,7 @@ class TestHybridRouteLayer: def test_initialization(self, openai_encoder, routes): route_layer = HybridRouteLayer(encoder=openai_encoder, routes=routes) assert route_layer.index is not None and route_layer.categories is not None + assert openai_encoder.score_threshold == 0.82 assert route_layer.score_threshold == 0.82 assert len(route_layer.index) == 5 assert len(set(route_layer.categories)) == 2 @@ -112,7 +113,8 @@ def test_pass_threshold(self, openai_encoder): def test_failover_score_threshold(self, base_encoder): route_layer = HybridRouteLayer(encoder=base_encoder) - assert route_layer.score_threshold == 0.82 + assert base_encoder.score_threshold == 0.50 + assert route_layer.score_threshold == 0.50 # Add more tests for edge cases and error handling as needed. From 20d16a8282825be530e032bc7cf40de9bf3287d0 Mon Sep 17 00:00:00 2001 From: Ismail Ashraq Date: Sun, 7 Jan 2024 19:16:19 +0500 Subject: [PATCH 3/7] remove test code --- semantic_router/encoders/bm25.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/semantic_router/encoders/bm25.py b/semantic_router/encoders/bm25.py index 69ca58ec..68150cb7 100644 --- a/semantic_router/encoders/bm25.py +++ b/semantic_router/encoders/bm25.py @@ -19,11 +19,7 @@ def __init__(self, name: str = "bm25", score_threshold: float = 0.82): "You can install it with: `pip install semantic-router[hybrid]`" ) logger.info("Downloading and initializing BM25 model parameters.") - # self.model = encoder.default() - self.model = encoder() - self.model.fit( - corpus=["test test", "this is another message", "hello how are you"] - ) + self.model = encoder.default() params = self.model.get_params() doc_freq = params["doc_freq"] From f67b051c84b8ca343ba1c749815867f106a0dc7a Mon Sep 17 00:00:00 2001 From: James Briggs <35938317+jamescalam@users.noreply.github.com> Date: Sun, 7 Jan 2024 16:19:17 +0100 Subject: [PATCH 4/7] add type definition for self.encoder --- semantic_router/layer.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/semantic_router/layer.py b/semantic_router/layer.py index dd9006ca..08261756 100644 --- a/semantic_router/layer.py +++ b/semantic_router/layer.py @@ -4,10 +4,7 @@ import numpy as np import yaml -from semantic_router.encoders import ( - BaseEncoder, - OpenAIEncoder -) +from semantic_router.encoders import BaseEncoder, OpenAIEncoder from semantic_router.linear import similarity_matrix, top_scores from semantic_router.llms import BaseLLM, OpenAILLM from semantic_router.route import Route @@ -153,6 +150,7 @@ class RouteLayer: index: np.ndarray | None = None categories: np.ndarray | None = None score_threshold: float + encoder: BaseEncoder def __init__( self, From 34135011c8770ac82b3a7c635f00e00f615e9934 Mon Sep 17 00:00:00 2001 From: James Briggs <35938317+jamescalam@users.noreply.github.com> Date: Sun, 7 Jan 2024 16:30:14 +0100 Subject: [PATCH 5/7] add test for encoder is None --- tests/unit/test_layer.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unit/test_layer.py b/tests/unit/test_layer.py index 286e7ace..92577f34 100644 --- a/tests/unit/test_layer.py +++ b/tests/unit/test_layer.py @@ -119,6 +119,10 @@ def test_initialization_different_encoders(self, cohere_encoder, openai_encoder) route_layer_openai = RouteLayer(encoder=openai_encoder) assert route_layer_openai.score_threshold == 0.82 + def test_initialization_no_encoder(self, openai_encoder): + route_layer_none = RouteLayer(encoder=None) + assert route_layer_none.score_threshold == openai_encoder.score_threshold + def test_initialization_dynamic_route(self, cohere_encoder, openai_encoder): route_layer_cohere = RouteLayer(encoder=cohere_encoder) assert route_layer_cohere.score_threshold == 0.3 From 89029f0177d10899da9c27909bd0fd05a5447f4f Mon Sep 17 00:00:00 2001 From: James Briggs <35938317+jamescalam@users.noreply.github.com> Date: Sun, 7 Jan 2024 16:48:30 +0100 Subject: [PATCH 6/7] add patch for encoder in None test --- tests/unit/test_layer.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_layer.py b/tests/unit/test_layer.py index 92577f34..cdf86bbc 100644 --- a/tests/unit/test_layer.py +++ b/tests/unit/test_layer.py @@ -120,8 +120,9 @@ def test_initialization_different_encoders(self, cohere_encoder, openai_encoder) assert route_layer_openai.score_threshold == 0.82 def test_initialization_no_encoder(self, openai_encoder): - route_layer_none = RouteLayer(encoder=None) - assert route_layer_none.score_threshold == openai_encoder.score_threshold + with patch("semantic_router.encoders.OpenAIEncoder") as _: + route_layer_none = RouteLayer(encoder=None) + assert route_layer_none.score_threshold == openai_encoder.score_threshold def test_initialization_dynamic_route(self, cohere_encoder, openai_encoder): route_layer_cohere = RouteLayer(encoder=cohere_encoder) From 534d9eef02acdc085b693c38fc32683c77ea22ed Mon Sep 17 00:00:00 2001 From: James Briggs <35938317+jamescalam@users.noreply.github.com> Date: Sun, 7 Jan 2024 16:59:27 +0100 Subject: [PATCH 7/7] add env var --- tests/unit/test_layer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_layer.py b/tests/unit/test_layer.py index cdf86bbc..6e2ac317 100644 --- a/tests/unit/test_layer.py +++ b/tests/unit/test_layer.py @@ -120,9 +120,9 @@ def test_initialization_different_encoders(self, cohere_encoder, openai_encoder) assert route_layer_openai.score_threshold == 0.82 def test_initialization_no_encoder(self, openai_encoder): - with patch("semantic_router.encoders.OpenAIEncoder") as _: - route_layer_none = RouteLayer(encoder=None) - assert route_layer_none.score_threshold == openai_encoder.score_threshold + os.environ["OPENAI_API_KEY"] = "test_api_key" + route_layer_none = RouteLayer(encoder=None) + assert route_layer_none.score_threshold == openai_encoder.score_threshold def test_initialization_dynamic_route(self, cohere_encoder, openai_encoder): route_layer_cohere = RouteLayer(encoder=cohere_encoder)