diff --git a/filters/toxicity/filter.py b/filters/toxicity/filter.py index f958e8c54..9d6210976 100644 --- a/filters/toxicity/filter.py +++ b/filters/toxicity/filter.py @@ -18,16 +18,30 @@ class ToxicityTypes(str, enum.Enum): class ToxicityFilter(SentenceOperation): + + """ + Filter to select sentences based on their toxicity. + """ + tasks = [TaskType.TEXT_TO_TEXT_GENERATION] languages = ["en"] + keywords = ["semantic", "social-reasoning"] heavy = True def __init__( self, - toxicity_type: ToxicityTypes, + toxicity_type: ToxicityTypes = ToxicityTypes.TOXICITY, op: str = ">", threshold: float = 0.5, ): + """ + Initialize filter by specifying the type and threshold values for toxicity. + :param toxicity_type: One of ToxicityTypes.TOXICITY, ToxicityTypes.SEVERE_TOXICITY, + ToxicityTypes.OBSCENE, ToxicityTypes.IDENTITY_ATTACK, + ToxicityTypes.INSULT, ToxicityTypes.THREAT, ToxicityTypes.SEXUAL_EXPLICIT + :param op: Operator to compare the toxicity value to the threshold. One of ">", "<", ">=", "<=", "==". + :param threshold: A float value between 0.0 and 1.0 for comparing toxicity of an input text. + """ super().__init__() self.check_threshold_value(threshold)