🎨 reformatted actions.py

Author: m-misiura

nemoguardrails/library/embedding_topic_detector/actions.py

@@ -27,25 +27,46 @@
 
 
 class EmbeddingTopicDetector:
-    def __init__(self, embedding_model: str, embedding_engine: str, examples: Dict[str, List[str]], threshold: float, top_k: int):
+    def __init__(
+        self,
+        embedding_model: str,
+        embedding_engine: str,
+        examples: Dict[str, List[str]],
+        threshold: float,
+        top_k: int,
+    ):
         self.threshold = threshold
         self.top_k = top_k
         self.model = init_embedding_model(embedding_model, embedding_engine)
         self.embeddings = {
             cat: [np.array(e) for e in self.model.encode(queries)]
-            for cat, queries in examples.items() if queries
+            for cat, queries in examples.items()
+            if queries
         }
 
     async def detect(self, query: str) -> Dict:
         query_emb = np.array((await self.model.encode_async([query]))[0])
 
         sims = sorted(
-            [(cat, float(np.dot(query_emb, emb) / (np.linalg.norm(query_emb) * np.linalg.norm(emb) or 1)))
-             for cat, embs in self.embeddings.items() for emb in embs],
-            key=lambda x: x[1], reverse=True
+            [
+                (
+                    cat,
+                    float(
+                        np.dot(query_emb, emb)
+                        / (np.linalg.norm(query_emb) * np.linalg.norm(emb) or 1)
+                    ),
+                )
+                for cat, embs in self.embeddings.items()
+                for emb in embs
+            ],
+            key=lambda x: x[1],
+            reverse=True,
         )[: self.top_k]
 
-        scores = {cat: float(np.mean([s for c, s in sims if c == cat]) or 0.0) for cat in self.embeddings}
+        scores = {
+            cat: float(np.mean([s for c, s in sims if c == cat]) or 0.0)
+            for cat in self.embeddings
+        }
         max_score = max(scores.values(), default=0.0)
 
         return {
@@ -62,22 +83,34 @@ async def _check(context: Optional[dict], llm_task_manager, message_key: str) ->
 
     if cache_key not in _detector_cache:
         _detector_cache[cache_key] = EmbeddingTopicDetector(
-            config["embedding_model"], config["embedding_engine"],
-            config["examples"], config.get("threshold", 0.75), config.get("top_k", 3)
+            config["embedding_model"],
+            config["embedding_engine"],
+            config["examples"],
+            config.get("threshold", 0.75),
+            config.get("top_k", 3),
         )
 
     query = context.get(message_key) if context else None
     if not query:
-        return {"on_topic": True, "confidence": 0.0, "top_category": None, "category_scores": {}}
+        return {
+            "on_topic": True,
+            "confidence": 0.0,
+            "top_category": None,
+            "category_scores": {},
+        }
 
     return await _detector_cache[cache_key].detect(query)
 
 
 @action(is_system_action=True)
-async def embedding_topic_check(context: Optional[dict] = None, llm_task_manager=None) -> dict:
+async def embedding_topic_check(
+    context: Optional[dict] = None, llm_task_manager=None
+) -> dict:
     return await _check(context, llm_task_manager, "user_message")
 
 
 @action(is_system_action=True)
-async def embedding_topic_check_output(context: Optional[dict] = None, llm_task_manager=None) -> dict:
+async def embedding_topic_check_output(
+    context: Optional[dict] = None, llm_task_manager=None
+) -> dict:
     return await _check(context, llm_task_manager, "bot_message")