Spaces:
Running
Running
| """Scoring server for Toxicity-Lens, compatible with the Perspective API. | |
| POST /v1alpha1/comments:analyze the Perspective API's request and response JSON (also served at /v1/comments:analyze) | |
| POST /v1alpha1/comments:suggestscore accepted so clients that send feedback keep working; stored only if FEEDBACK_FILE is set | |
| GET /$discovery/rest discovery document, so Google API client libraries can be pointed at this server | |
| POST /v1/moderations OpenAI-style moderation shape, from the same scores | |
| POST /beta/litellm_basic_guardrail_api LiteLLM's Generic Guardrail API (guardrail: generic_guardrail_api, api_base: this server) | |
| GET /health | |
| GPU or CPU (PyTorch): MODEL_DIR=abacusai/Toxicity-Lens uvicorn server:app --host 0.0.0.0 --port 8080 | |
| CPU without PyTorch: MODEL_DIR=abacusai/Toxicity-Lens BACKEND=onnx uvicorn server:app --host 0.0.0.0 --port 8080 | |
| Optional: ONNX_FILE (default onnx/model_int8_block32.onnx, 8-bit, about 0.001 from float32 on average on any CPU; onnx/model.onnx is | |
| the exact float32 export; onnx/model_int8.onnx is plain dynamic int8 for transformers.js), API_KEYS=k1,k2 (requests must | |
| carry one as ?key=, x-goog-api-key or a Bearer token), PUBLIC_URL (rootUrl in the discovery document, for use behind a proxy), | |
| MODEL_REVISION (a commit hash of the model repo: pinning it keeps every score reproducible), ORT_THREADS (caps onnxruntime's | |
| threads; set it to the container's CPU count, for example 2 on a Hugging Face cpu-basic Space). | |
| """ | |
| import json | |
| import os | |
| import re | |
| import threading | |
| import uuid | |
| import numpy as np | |
| from fastapi import FastAPI, Request | |
| from fastapi.concurrency import run_in_threadpool | |
| from fastapi.responses import JSONResponse | |
| ATTRS = ["TOXICITY", "SEVERE_TOXICITY", "IDENTITY_ATTACK", "INSULT", "PROFANITY", "THREAT", "SEXUALLY_EXPLICIT"] | |
| MAX_BYTES = 20480 # Perspective's documented per-comment limit | |
| MAX_TOKENS, STRIDE, BATCH = 256, 32, 32 # training, calibration and every evaluation scored 256-token windows | |
| SENTENCE = re.compile(r"\S[^.!?。！？\n]*[.!?。！？]*") # spanAnnotations: Perspective returns per-sentence spans | |
| # OpenAI moderation categories we can fill honestly from our attributes (plus our own "toxicity"). | |
| OPENAI_MAP = {"harassment": "INSULT", "harassment/threatening": "THREAT", "hate": "IDENTITY_ATTACK", | |
| "sexual": "SEXUALLY_EXPLICIT", "violence": "THREAT", "toxicity": "TOXICITY"} | |
| class Scorer: | |
| """Scores texts with a saved Toxicity-Lens model. A text longer than one 256-token window is scored in overlapping | |
| windows and each attribute keeps its highest window score, so an insult deep inside a long comment still counts.""" | |
| def __init__(self, path: str, backend: str = "torch", onnx_file: str = "onnx/model_int8_block32.onnx", logits_fn=None, revision=None): | |
| from tokenizers import Tokenizer | |
| if not os.path.isdir(path): | |
| from huggingface_hub import snapshot_download | |
| weights = ["model.safetensors"] if backend == "torch" else [onnx_file, onnx_file + ".data"] | |
| path = snapshot_download(path, revision=revision, allow_patterns=["config.json", "tokenizer.json", "tokenizer_config.json", *weights]) | |
| labels = json.load(open(f"{path}/config.json"))["id2label"] | |
| assert [labels[str(i)] for i in range(len(labels))] == ATTRS, f"model labels {labels} are not the Perspective attribute set" | |
| self.tok = Tokenizer.from_file(f"{path}/tokenizer.json") | |
| self.tok.enable_truncation(MAX_TOKENS, stride=STRIDE) | |
| self.tok.no_padding() | |
| self.pad_id = self.tok.token_to_id(json.load(open(f"{path}/tokenizer_config.json"))["pad_token"]) | |
| self.lock = threading.Lock() # ponytail: one forward pass at a time; add micro-batching if one GPU must serve >200 requests/s | |
| self.logits = logits_fn or (self._onnx(path, onnx_file) if backend == "onnx" else self._torch(path)) | |
| def _onnx(path, onnx_file): | |
| import onnxruntime as ort | |
| so = ort.SessionOptions() | |
| so.intra_op_num_threads = int(os.environ.get("ORT_THREADS", "0")) # 0 = onnxruntime's default, one thread per core it detects | |
| sess = ort.InferenceSession(f"{path}/{onnx_file}", so, providers=["CPUExecutionProvider"]) | |
| return lambda ids, mask: sess.run(None, {"input_ids": ids, "attention_mask": mask})[0] | |
| def _torch(path): | |
| import torch | |
| from transformers import AutoModelForSequenceClassification | |
| dev = "cuda" if torch.cuda.is_available() else "cpu" | |
| dtype = torch.bfloat16 if dev == "cuda" else torch.float32 | |
| model = AutoModelForSequenceClassification.from_pretrained(path, dtype=dtype).to(dev).eval() | |
| def run(ids, mask): | |
| with torch.no_grad(): | |
| out = model(input_ids=torch.from_numpy(ids).to(dev), attention_mask=torch.from_numpy(mask).to(dev)) | |
| return out.logits.float().cpu().numpy() | |
| return run | |
| def __call__(self, texts): | |
| """One (summary {attribute: probability}, [(begin, end, {attribute: probability}) per window]) per text.""" | |
| wins, owner = [], [] | |
| for i, enc in enumerate(self.tok.encode_batch(list(texts))): | |
| for w in [enc, *enc.overflowing]: | |
| wins.append(w) | |
| owner.append(i) | |
| probs = [] | |
| with self.lock: | |
| for b in range(0, len(wins), BATCH): | |
| part = wins[b:b + BATCH] | |
| ids = np.full((len(part), max(len(w.ids) for w in part)), self.pad_id, dtype=np.int64) | |
| mask = np.zeros_like(ids) | |
| for j, w in enumerate(part): | |
| ids[j, :len(w.ids)], mask[j, :len(w.ids)] = w.ids, 1 | |
| probs.append(1 / (1 + np.exp(-self.logits(ids, mask).astype(np.float64)))) | |
| out = [[None, []] for _ in texts] | |
| for w, i, p in zip(wins, owner, np.concatenate(probs)): | |
| offs = [o for o, special in zip(w.offsets, w.special_tokens_mask) if not special] or [(0, 0)] | |
| out[i][1].append((offs[0][0], offs[-1][1], dict(zip(ATTRS, p.tolist())))) | |
| out[i][0] = p if out[i][0] is None else np.maximum(out[i][0], p) | |
| return [(dict(zip(ATTRS, s.tolist())), spans) for s, spans in out] | |
| def error(code: int, message: str, status: str = "INVALID_ARGUMENT"): | |
| return JSONResponse(status_code=code, content={"error": {"code": code, "message": message, "status": status}}) | |
| def discovery_doc(root: str) -> dict: | |
| """Just enough of a discovery document for googleapiclient / googleapis to build a commentanalyzer v1alpha1 client.""" | |
| def method(name, request, response): | |
| return {"id": f"commentanalyzer.comments.{name}", "path": f"v1alpha1/comments:{name}", "flatPath": f"v1alpha1/comments:{name}", | |
| "httpMethod": "POST", "parameters": {}, "parameterOrder": [], "request": {"$ref": request}, "response": {"$ref": response}} | |
| names = ["AnalyzeCommentRequest", "AnalyzeCommentResponse", "SuggestCommentScoreRequest", "SuggestCommentScoreResponse"] | |
| return {"kind": "discovery#restDescription", "discoveryVersion": "v1", "id": "commentanalyzer:v1alpha1", "name": "commentanalyzer", | |
| "version": "v1alpha1", "title": "Toxicity-Lens (compatible with the Perspective API)", "protocol": "rest", | |
| "rootUrl": root, "servicePath": "", "baseUrl": root, "batchPath": "batch", | |
| "parameters": {"key": {"type": "string", "location": "query"}, "alt": {"type": "string", "location": "query", "default": "json"}}, | |
| "schemas": {n: {"id": n, "type": "object"} for n in names}, | |
| "resources": {"comments": {"methods": {"analyze": method("analyze", *names[:2]), "suggestscore": method("suggestscore", *names[2:])}}}} | |
| def create_app(scorer, model_name: str = "toxicity-lens") -> FastAPI: | |
| app = FastAPI(title="Toxicity-Lens scorer, compatible with the Perspective API") | |
| keys = {k for k in os.environ.get("API_KEYS", "").split(",") if k} | |
| feedback = os.environ.get("FEEDBACK_FILE") | |
| def denied(request: Request): | |
| if not keys: | |
| return None | |
| auth = request.headers.get("authorization", "") | |
| key = request.query_params.get("key") or request.headers.get("x-goog-api-key") or (auth[7:] if auth.lower().startswith("bearer ") else None) | |
| if not key: | |
| return error(403, "The request is missing a valid API key.", "PERMISSION_DENIED") | |
| return None if key in keys else error(400, "API key not valid. Please pass a valid API key.") | |
| async def read_json(request: Request): | |
| try: | |
| body = await request.json() | |
| except Exception: | |
| return None | |
| return body if isinstance(body, dict) else None | |
| def health(): | |
| return {"status": "ok", "model": model_name} | |
| def discovery(request: Request): | |
| return discovery_doc(os.environ.get("PUBLIC_URL", str(request.base_url)).rstrip("/") + "/") | |
| async def analyze(request: Request): | |
| if (refused := denied(request)) is not None: | |
| return refused | |
| body = await read_json(request) | |
| if body is None: | |
| return error(400, "Invalid JSON payload received.") | |
| comment = body.get("comment") | |
| text = comment.get("text") if isinstance(comment, dict) else None | |
| if not isinstance(text, str) or not text.strip(): | |
| return error(400, "Comment must be non-empty.") | |
| if comment.get("type", "PLAIN_TEXT") != "PLAIN_TEXT": | |
| return error(400, "Only PLAIN_TEXT comments are supported.") | |
| if len(text.encode("utf-8")) > MAX_BYTES: | |
| return error(400, f"Comment text too long. Must be at most {MAX_BYTES} bytes.") | |
| requested = body.get("requestedAttributes") | |
| if not isinstance(requested, dict) or not requested: | |
| return error(400, "Must specify at least one requestedAttribute.") | |
| for attr, opts in requested.items(): | |
| if attr not in ATTRS: | |
| return error(400, f"Attribute {attr} is not supported by this server.") | |
| if isinstance(opts, dict) and opts.get("scoreType", "PROBABILITY") != "PROBABILITY": | |
| return error(400, f"Score type {opts['scoreType']} is not supported.") | |
| summary = (await run_in_threadpool(scorer, [text]))[0][0] | |
| spans = [(0, len(text), summary)] | |
| if body.get("spanAnnotations"): | |
| sents = [m.span() for m in SENTENCE.finditer(text)] | |
| spans = [(b, e, s) for (b, e), (s, _) in zip(sents, await run_in_threadpool(scorer, [text[b:e] for b, e in sents]))] | |
| out = {} | |
| for attr, opts in requested.items(): | |
| threshold = opts.get("scoreThreshold") if isinstance(opts, dict) else None | |
| if threshold is not None and summary[attr] < float(threshold): | |
| continue # Perspective omits attributes scored below the requested threshold | |
| out[attr] = {"spanScores": [{"begin": b, "end": e, "score": {"value": s[attr], "type": "PROBABILITY"}} for b, e, s in spans], | |
| "summaryScore": {"value": summary[attr], "type": "PROBABILITY"}} | |
| resp = {"attributeScores": out, "languages": body.get("languages") or ["und"]} # no language detection: echoes the request | |
| if "clientToken" in body: | |
| resp["clientToken"] = body["clientToken"] | |
| return resp | |
| app.post("/v1alpha1/comments:analyze")(analyze) | |
| app.post("/v1/comments:analyze")(analyze) | |
| async def suggest_score(request: Request): | |
| if (refused := denied(request)) is not None: | |
| return refused | |
| body = await read_json(request) | |
| if body is None or not isinstance(body.get("comment"), dict) or not isinstance(body.get("attributeScores"), dict): | |
| return error(400, "Request must include comment and attributeScores.") | |
| if feedback: | |
| with open(feedback, "a", encoding="utf-8") as f: | |
| f.write(json.dumps(body, ensure_ascii=False) + "\n") | |
| resp = {"requestedLanguages": body.get("languages") or []} | |
| if "clientToken" in body: | |
| resp["clientToken"] = body["clientToken"] | |
| return resp | |
| async def moderations(request: Request): | |
| if (refused := denied(request)) is not None: | |
| return refused | |
| body = await read_json(request) | |
| inp = body.get("input") if body else None | |
| texts = [inp] if isinstance(inp, str) else inp | |
| if not isinstance(texts, list) or not texts or not all(isinstance(t, str) and t for t in texts): | |
| return JSONResponse(status_code=400, content={"error": {"message": "input must be a non-empty string or list of strings", | |
| "type": "invalid_request_error"}}) | |
| results = [] | |
| for summary, _ in await run_in_threadpool(scorer, texts): | |
| cs = {cat: summary[attr] for cat, attr in OPENAI_MAP.items()} | |
| results.append({"flagged": any(v >= 0.5 for v in cs.values()), "categories": {k: v >= 0.5 for k, v in cs.items()}, | |
| "category_scores": cs}) | |
| return {"id": f"modr-{uuid.uuid4().hex[:24]}", "model": body.get("model") or model_name, "results": results} | |
| async def litellm_guardrail(request: Request): | |
| """LiteLLM's Generic Guardrail API: BLOCKED when any text reaches the threshold on any chosen attribute. Set | |
| additional_provider_specific_params {threshold, attributes} in the LiteLLM guardrail config to change the defaults.""" | |
| if (refused := denied(request)) is not None: | |
| return refused | |
| body = await read_json(request) or {} | |
| texts = [t for t in body.get("texts") or [] if isinstance(t, str) and t.strip()] | |
| params = body.get("additional_provider_specific_params") or {} | |
| threshold = float(params.get("threshold", 0.8)) | |
| attrs = [a for a in params.get("attributes") or ATTRS if a in ATTRS] | |
| hits = sorted({a for s, _ in await run_in_threadpool(scorer, texts) for a in attrs if s[a] >= threshold}) if texts else [] | |
| if hits: | |
| return {"action": "BLOCKED", "blocked_reason": f"{model_name}: {', '.join(hits)} at or above {threshold}"} | |
| return {"action": "NONE"} | |
| return app | |
| if os.environ.get("MODEL_DIR"): | |
| app = create_app(Scorer(os.environ["MODEL_DIR"], os.environ.get("BACKEND", "torch"), os.environ.get("ONNX_FILE", "onnx/model_int8_block32.onnx"), | |
| revision=os.environ.get("MODEL_REVISION")), os.environ.get("MODEL_NAME", "toxicity-lens")) | |