Added type hints and hedging lang in to_dict

arthur-ai · Nov 17, 2023 · a327025 · a327025
1 parent 1722013
commit a327025
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/arthur_bench/scoring/hedging_language.py b/arthur_bench/scoring/hedging_language.py
@@ -19,7 +19,9 @@ class HedgingLanguage(Scorer):
     model output.
     """
 
-    def __init__(self, model_type=DEFAULT_MODEL, hedging_language=DEFAULT_HEDGE):
+    def __init__(
+        self, model_type: str = DEFAULT_MODEL, hedging_language: str = DEFAULT_HEDGE
+    ):
         """
         Hedging Language score implementation.
 
@@ -40,7 +42,10 @@ def requires_reference() -> bool:
         return False
 
     def to_dict(self, warn=False):
-        return {"model_type": self.scorer.model_type}
+        return {
+            "model_type": self.scorer.model_type,
+            "hedging_language": self.hedging_language,
+        }
 
     def run_batch(
         self,