yetessam commited on
Commit
7c011b8
·
verified ·
1 Parent(s): db26c1f

Update tools/polite_guard.py

Browse files
Files changed (1) hide show
  1. tools/polite_guard.py +5 -2
tools/polite_guard.py CHANGED
@@ -32,7 +32,7 @@ class PoliteGuardTool(Tool):
32
  self.score = None
33
 
34
  #@tool
35
- def ask_polite_guard(self, input_text: str) -> str:
36
  """
37
 
38
  Args:
@@ -43,7 +43,10 @@ class PoliteGuardTool(Tool):
43
  """
44
  try:
45
  classifier = pipeline("text-classification", "Intel/polite-guard")
46
- return classifier(input_text)
 
 
 
47
 
48
  except Exception as e:
49
  return f"Error fetching classification for text '{input_text}': {str(e)}"
 
32
  self.score = None
33
 
34
  #@tool
35
+ def ask_polite_guard(self, input_text: str) -> tuple[str, float]:
36
  """
37
 
38
  Args:
 
43
  """
44
  try:
45
  classifier = pipeline("text-classification", "Intel/polite-guard")
46
+ result = classifier(input_text)[0]
47
+ print(result)
48
+ print(len(result))
49
+ return result['label'], result['score']
50
 
51
  except Exception as e:
52
  return f"Error fetching classification for text '{input_text}': {str(e)}"