Granite guardian minor changes (#1605)

* Granite guardian: move tokenizer init to prepare() Signed-off-by: Martín Santillán Cooper <[email protected]> * Add Custom risk entry to RISK_TYPE_TO_CLASS Signed-off-by: Martín Santillán Cooper <[email protected]> --------- Signed-off-by: Martín Santillán Cooper <[email protected]> Co-authored-by: Elron Bandel <[email protected]>
IBM · Feb 18, 2025 · d6c0c50 · d6c0c50
1 parent 3dd1f9c
commit d6c0c50
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/src/unitxt/metrics.py b/src/unitxt/metrics.py
@@ -5921,8 +5921,11 @@ class GraniteGuardianBase(InstanceMetric):
     _requirements_list: List[str] = ["torch", "transformers"]
 
     def prepare(self):
+        from transformers import AutoTokenizer
         if not isinstance(self.risk_type, RiskType):
             self.risk_type = RiskType[self.risk_type]
+        if not hasattr(self, "_tokenizer") or self._tokenizer is None:
+            self._tokenizer = AutoTokenizer.from_pretrained(self.hf_model_name)
 
     def verify(self):
         super().verify()
@@ -5957,12 +5960,9 @@ def get_prompt(self, messages):
         )
 
     def compute(self, references: List[Any], prediction: Any, task_data: Dict) -> dict:
-        from transformers import AutoTokenizer
-
         self.verify_granite_guardian_config(task_data)
         self.set_main_score()
-        if not hasattr(self, "_tokenizer") or self._tokenizer is None:
-            self._tokenizer = AutoTokenizer.from_pretrained(self.hf_model_name)
+
         if self.inference_engine is None:
             self.inference_engine = WMLInferenceEngineGeneration(
                 model_name=self.wml_model_name,
@@ -6178,6 +6178,7 @@ def process_input_fields(self, task_data):
     RiskType.ASSISTANT_MESSAGE: GraniteGuardianAssistantRisk,
     RiskType.RAG: GraniteGuardianRagRisk,
     RiskType.AGENTIC: GraniteGuardianAgenticRisk,
+    RiskType.CUSTOM_RISK: GraniteGuardianCustomRisk,
 }
 
 class ExecutionAccuracy(InstanceMetric):