| from enum import Enum |
| from typing import List |
| import re |
|
|
| class ToxicityLevel(Enum): |
| SAFE = "safe" |
| MILD = "mild" |
| EXPLICIT = "explicit" |
| SLUR = "slur" |
| THREAT = "threat" |
|
|
| class ContextClassifier: |
| def __init__(self): |
| |
| self._mild = {'damn', 'crap', 'hell', 'ass'} |
| self._explicit = {'fuck', 'shit', 'bitch', 'piss', 'dick', 'cock', 'pussy'} |
| self._slurs = {'bastard'} |
| self._threat_keywords = ['kill', 'die', 'death', 'hurt', 'harm'] |
| |
| def classify_context(self, text: str) -> ToxicityLevel: |
| """ |
| Classify the toxicity level of text with context awareness. |
| |
| Args: |
| text: Input text to classify |
| |
| Returns: |
| ToxicityLevel: The classified toxicity level |
| """ |
| text_lower = text.lower() |
| words = set(re.findall(r'\b\w+\b', text_lower)) |
| |
| |
| if any(keyword in text_lower for keyword in self._threat_keywords): |
| return ToxicityLevel.THREAT |
| |
| |
| if words.intersection(self._slurs): |
| return ToxicityLevel.SLUR |
| |
| |
| if words.intersection(self._explicit): |
| return ToxicityLevel.EXPLICIT |
| |
| |
| if words.intersection(self._mild): |
| return ToxicityLevel.MILD |
| |
| return ToxicityLevel.SAFE |
