Add Custom risk entry to RISK_TYPE_TO_CLASS

Signed-off-by: Martín Santillán Cooper <[email protected]>
IBM · Feb 12, 2025 · 75f99b0 · 75f99b0
1 parent c0956ef
commit 75f99b0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/unitxt/metrics.py b/src/unitxt/metrics.py
@@ -6175,6 +6175,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):