Fix default behavior toxic_label in HuggingFaceToxicityModel. (#1146)

evidentlyai · Jun 6, 2024 · e68824a · e68824a
1 parent 89f06c6
commit e68824a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/evidently/features/hf_feature.py b/src/evidently/features/hf_feature.py
@@ -94,7 +94,10 @@ def _toxicity(model_name: Optional[str], toxic_label: Optional[str], data: pd.Se
 
     column_data = data.values.tolist()
     model = evaluate.load("toxicity", model_name, module_type="measurement")
-    scores = model.compute(predictions=column_data, toxic_label=toxic_label)
+    if toxic_label is None:
+        scores = model.compute(predictions=column_data)
+    else:
+        scores = model.compute(predictions=column_data, toxic_label=toxic_label)
     return pd.Series(scores["toxicity"], index=data.index)