Spaces:
Sleeping
Sleeping
File size: 2,423 Bytes
9ee432e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
# tools/profanity_guard.py
from typing import Any, Dict
from smolagents.tools import Tool
import json
class ProfanityGuardTool(Tool):
name = "profanity_guard"
description = "Detects profanity in English text and returns a label and confidence."
inputs: Dict[str, Dict[str, Any]] = {
"text": {"type": "string", "description": "English text to check for profanity."}
}
output_type = "string" # return JSON string to match your web_search.py pattern
def __init__(self, model_name: str = "tarekziade/pardonmyai", device: int | None = None, **kwargs: Any) -> None:
"""
model_name options:
- "tarekziade/pardonmyai" (default, DistilBERT-based, binary PROFANE/CLEAN)
- "tarekziade/pardonmyai-tiny" (smaller, faster)
"""
super().__init__()
try:
import torch # noqa: F401
from transformers import pipeline # type: ignore
except ImportError as e:
raise ImportError(
"You must install `transformers` (and optionally `torch`) to use ProfanityGuardTool.\n"
"Example: pip install transformers torch --extra-index-url https://download.pytorch.org/whl/cu121"
) from e
self.model_name = model_name
# Pick device automatically if not specified
try:
import torch
if device is None:
device = 0 if torch.cuda.is_available() else -1
except Exception:
device = -1 # CPU fallback if torch not available/working
# Build the pipeline once (fast subsequent calls)
from transformers import pipeline
self.pipe = pipeline(
task="sentiment-analysis", # model card uses this task name
model=self.model_name,
device=device,
truncation=True
)
def forward(self, text: str) -> str:
t = (text or "").strip()
if not t:
raise ValueError("`text` must be a non-empty string.")
# Light normalization so profanity isn't split by odd whitespace
t = " ".join(t.split())
out = self.pipe(t)[0] # e.g. {'label': 'PROFANE'|'CLEAN', 'score': 0.xx}
payload = {
"model": self.model_name,
"label": str(out.get("label", "")),
"score": float(out.get("score", 0.0)),
}
return json.dumps(payload)
|