Update refusal_shield.py

genia-dev · Dec 17, 2023 · 13480ab · 13480ab
1 parent b725e84
commit 13480ab
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/vibraniumdome-shields/vibraniumdome_shields/shields/output/refusal_shield.py b/vibraniumdome-shields/vibraniumdome_shields/shields/output/refusal_shield.py
@@ -1,5 +1,5 @@
 import logging
-from typing import List
+from typing import List, Optional
 from uuid import UUID
 import math
 from transformers import pipeline
@@ -8,9 +8,9 @@
 
 
 class RefusalShieldDeflectionResult(ShieldDeflectionResult):
-    model: str = ""
-    labels: List
-    scores: List
+    model: Optional[str] = None
+    labels: Optional[List] = None
+    scores: Optional[List] = None
 
 
 class RefusalShield(VibraniumShield):
@@ -50,6 +50,6 @@ def deflect(self, llm_interaction: LLMInteraction, shield_policy_config: dict, s
             if risk >= threshold:
                 shield_results.append(RefusalShieldDeflectionResult(model=self._model, labels=categories, scores=[risk, round(1 - risk, 2)], risk=risk))
             else:
-                shield_results.append(RefusalShieldDeflectionResult(model=self._model, labels=categories, scores=[risk, round(1 - risk, 2)], risk=0))
+                shield_results.append(RefusalShieldDeflectionResult())
 
         return shield_results