diff --git a/semantic_router/layer.py b/semantic_router/layer.py index dd746d0ec6ee73fa1d5db4ba533d37e3878b845a..fcedc979b94e2c07358ca4ff17be44705ff1d888 100644 --- a/semantic_router/layer.py +++ b/semantic_router/layer.py @@ -69,7 +69,9 @@ class DecisionLayer: scores, idx = top_scores(sim, top_k) # get the utterance categories (decision names) decisions = self.categories[idx] if self.categories is not None else [] - return [{"decision": d, "score": s.item()} for d, s in zip(decisions, scores)] + return [ + {"decision": d, "score": s.item()} for d, s in zip(decisions, scores) + ] else: return [] @@ -84,7 +86,9 @@ class DecisionLayer: scores_by_class[decision] = [score] # Calculate total score for each class - total_scores = {decision: sum(scores) for decision, scores in scores_by_class.items()} + total_scores = { + decision: sum(scores) for decision, scores in scores_by_class.items() + } top_class = max(total_scores, key=lambda x: total_scores[x], default=None) # Return the top class and its associated scores diff --git a/tests/functional/test_linear.py b/tests/functional/test_linear.py index 6771fd9ccd22f49b7af0e01339c10d9531da7218..210de6d200495cb37c3c46ab11c792efb8505dea 100644 --- a/tests/functional/test_linear.py +++ b/tests/functional/test_linear.py @@ -1,5 +1,5 @@ -import pytest import numpy as np +import pytest from semantic_router.linear import similarity_matrix, top_scores