diff --git a/semantic_router/hybrid_layer.py b/semantic_router/hybrid_layer.py index f293891a..5f223384 100644 --- a/semantic_router/hybrid_layer.py +++ b/semantic_router/hybrid_layer.py @@ -25,7 +25,7 @@ def __init__( routes: List[Route] = [], alpha: float = 0.3, top_k: int = 5, - aggregation: str = "SUM", + aggregation: str = "sum", ): self.encoder = encoder self.score_threshold = self.encoder.score_threshold @@ -41,7 +41,7 @@ def __init__( if self.top_k < 1: raise ValueError(f"top_k needs to be >= 1, but was: {self.top_k}.") self.aggregation = aggregation - if self.aggregation not in ["SUM", "MEAN", "MAX"]: + if self.aggregation not in ["sum", "mean", "max"]: raise ValueError( f"Unsupported aggregation method chosen: {aggregation}. Choose either 'SUM', 'MEAN', or 'MAX'." ) @@ -172,12 +172,12 @@ def _convex_scaling(self, dense: np.ndarray, sparse: np.ndarray): sparse = np.array(sparse) * (1 - self.alpha) return dense, sparse - def _set_aggregation_method(self, aggregation: str = "SUM"): - if aggregation == "SUM": + def _set_aggregation_method(self, aggregation: str = "sum"): + if aggregation == "sum": return lambda x: sum(x) - elif aggregation == "MEAN": + elif aggregation == "mean": return lambda x: np.mean(x) - elif aggregation == "MAX": + elif aggregation == "max": return lambda x: max(x) else: raise ValueError( diff --git a/semantic_router/layer.py b/semantic_router/layer.py index 9a856a86..221de2be 100644 --- a/semantic_router/layer.py +++ b/semantic_router/layer.py @@ -183,7 +183,7 @@ def __init__( routes: Optional[List[Route]] = None, index: Optional[BaseIndex] = None, # type: ignore top_k: int = 5, - aggregation: str = "SUM", + aggregation: str = "sum", ): logger.info("local") self.index: BaseIndex = index if index is not None else LocalIndex() @@ -202,7 +202,7 @@ def __init__( if self.top_k < 1: raise ValueError(f"top_k needs to be >= 1, but was: {self.top_k}.") self.aggregation = aggregation - if self.aggregation not in ["SUM", "MEAN", "MAX"]: + if self.aggregation not in ["sum", "mean", "max"]: raise ValueError( f"Unsupported aggregation method chosen: {aggregation}. Choose either 'SUM', 'MEAN', or 'MAX'." ) @@ -403,12 +403,12 @@ def _retrieve(self, xq: Any, top_k: int = 5) -> List[dict]: scores, routes = self.index.query(vector=xq, top_k=top_k) return [{"route": d, "score": s.item()} for d, s in zip(routes, scores)] - def _set_aggregation_method(self, aggregation: str = "SUM"): - if aggregation == "SUM": + def _set_aggregation_method(self, aggregation: str = "sum"): + if aggregation == "sum": return lambda x: sum(x) - elif aggregation == "MEAN": + elif aggregation == "mean": return lambda x: np.mean(x) - elif aggregation == "MAX": + elif aggregation == "max": return lambda x: max(x) else: raise ValueError( diff --git a/tests/unit/test_hybrid_layer.py b/tests/unit/test_hybrid_layer.py index 4b65b4cf..bf0c2ad2 100644 --- a/tests/unit/test_hybrid_layer.py +++ b/tests/unit/test_hybrid_layer.py @@ -194,7 +194,7 @@ def test_add_route_tfidf(self, cohere_encoder, tfidf_encoder, routes): assert len(hybrid_route_layer.sparse_index) == len(all_utterances) def test_setting_aggregation_methods(self, openai_encoder, routes): - for agg in ["SUM", "MEAN", "MAX"]: + for agg in ["sum", "mean", "max"]: route_layer = HybridRouteLayer( encoder=openai_encoder, sparse_encoder=sparse_encoder, @@ -217,7 +217,7 @@ def test_semantic_classify_multiple_routes_with_different_aggregation( {"route": "Route 3", "score": 0.1}, {"route": "Route 3", "score": 1.0}, ] - for agg in ["SUM", "MEAN", "MAX"]: + for agg in ["sum", "mean", "max"]: route_layer = HybridRouteLayer( encoder=openai_encoder, sparse_encoder=sparse_encoder, @@ -226,13 +226,13 @@ def test_semantic_classify_multiple_routes_with_different_aggregation( ) classification, score = route_layer._semantic_classify(route_scores) - if agg == "SUM": + if agg == "sum": assert classification == "Route 1" assert score == [0.5, 0.5, 0.5, 0.5] - elif agg == "MEAN": + elif agg == "mean": assert classification == "Route 2" assert score == [0.4, 0.6, 0.8] - elif agg == "MAX": + elif agg == "max": assert classification == "Route 3" assert score == [0.1, 1.0] diff --git a/tests/unit/test_layer.py b/tests/unit/test_layer.py index 06f24896..4a55777b 100644 --- a/tests/unit/test_layer.py +++ b/tests/unit/test_layer.py @@ -525,7 +525,7 @@ def test_remove(self): assert layer_config.routes == [] def test_setting_aggregation_methods(self, openai_encoder, routes): - for agg in ["SUM", "MEAN", "MAX"]: + for agg in ["sum", "mean", "max"]: route_layer = RouteLayer( encoder=openai_encoder, routes=routes, @@ -547,7 +547,7 @@ def test_semantic_classify_multiple_routes_with_different_aggregation( {"route": "Route 3", "score": 0.1}, {"route": "Route 3", "score": 1.0}, ] - for agg in ["SUM", "MEAN", "MAX"]: + for agg in ["sum", "mean", "max"]: route_layer = RouteLayer( encoder=openai_encoder, routes=routes, @@ -555,12 +555,12 @@ def test_semantic_classify_multiple_routes_with_different_aggregation( ) classification, score = route_layer._semantic_classify(route_scores) - if agg == "SUM": + if agg == "sum": assert classification == "Route 1" assert score == [0.5, 0.5, 0.5, 0.5] - elif agg == "MEAN": + elif agg == "mean": assert classification == "Route 2" assert score == [0.4, 0.6, 0.8] - elif agg == "MAX": + elif agg == "max": assert classification == "Route 3" assert score == [0.1, 1.0]