Update tools/polite_guard.py

tools/polite_guard.py

@@ -14,7 +14,7 @@ class PoliteGuardTool(Tool):
     output_type = "any"
 
     def forward(self, input_text: Any) -> Any:
-        self.label, self.score = self.ask_polite_guard(input_text)
+        self.label = self.ask_polite_guard(input_text)
         return self.label 
            
     def __init__(self, *args, **kwargs):
@@ -23,7 +23,7 @@ class PoliteGuardTool(Tool):
         self.score = None
         
     #@tool
-    def ask_polite_guard(self, input_text: str) -> dict:
+    def ask_polite_guard(self, input_text: str) -> str:
         """ 
     
         Args:
@@ -35,10 +35,7 @@ class PoliteGuardTool(Tool):
         try:
             classifier = pipeline("text-classification", "Intel/polite-guard")
             
-            return {
-                "label": classifier['label'],
-                "score": classifier['score']
-            }
+            return classifier['label']
             
         except Exception as e:
             return f"Error fetching classification for text '{input_text}': {str(e)}"