Blame for server.py

"""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))

    @staticmethod
    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]

    @staticmethod
    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

    @app.get("/health")
    def health():
        return {"status": "ok", "model": model_name}

    @app.get("/$discovery/rest")
    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)

    @app.post("/v1alpha1/comments:suggestscore")
    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

    @app.post("/v1/moderations")
    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}

    @app.post("/beta/litellm_basic_guardrail_api")
    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"))