Spaces:
Sleeping
Sleeping
File size: 3,657 Bytes
1048d56 | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | import json
import numpy as np, onnxruntime as ort, gradio as gr
from fastapi import FastAPI
from pydantic import BaseModel, Field
from huggingface_hub import snapshot_download
from tokenizers import Tokenizer
REPO, ONNX = "misukisu/moderation-modernbert-en", "onnx/model.onnx"
d = snapshot_download(REPO, allow_patterns=[ONNX, "tokenizer.json", "config.json", "thresholds.json"])
tok = Tokenizer.from_file(f"{d}/tokenizer.json")
tok.enable_truncation(256)
tok.enable_padding()
so = ort.SessionOptions()
so.intra_op_num_threads = 2
sess = ort.InferenceSession(f"{d}/{ONNX}", so, providers=["CPUExecutionProvider"])
id2label = json.load(open(f"{d}/config.json"))["id2label"]
LABELS = [id2label[str(i)] for i in range(len(id2label))]
TH = json.load(open(f"{d}/thresholds.json"))
MAX_BATCH = 32
def _run(texts):
texts = [str(t)[:4000] for t in texts][:MAX_BATCH]
if not texts:
return []
enc = tok.encode_batch(texts)
ids = np.array([e.ids for e in enc], dtype=np.int64)
mask = np.array([e.attention_mask for e in enc], dtype=np.int64)
probs = sess.run(None, {"input_ids": ids, "attention_mask": mask})[0]
return [{"flagged": bool(any(p[i] >= TH[l] for i, l in enumerate(LABELS))),
"categories": [l for i, l in enumerate(LABELS) if p[i] >= TH[l]],
"scores": {l: round(float(p[i]), 4) for i, l in enumerate(LABELS)}} for p in probs]
def moderate(text: str) -> dict:
"""Moderate one English text. Returns flagged, categories and per-label scores."""
return _run([text or ""])[0]
def moderate_batch(texts: list) -> list:
"""Moderate up to 32 texts at once."""
return _run(texts or [])
def ui(text):
r = moderate(text)
head = "### Flagged: " + ", ".join(r["categories"]) if r["flagged"] else "### Clean"
return head, {l: s for l, s in r["scores"].items()}, r
with gr.Blocks(title="Moderation API") as demo:
gr.Markdown("# Moderation API\nEnglish content moderation with ModernBERT, 13 labels. "
"Free to use from code, see the **Use via API** link at the bottom. "
f"Model: [{REPO}](https://huggingface.co/{REPO})")
with gr.Row():
with gr.Column():
inp = gr.Textbox(label="Text", lines=5, placeholder="Paste a message...")
btn = gr.Button("Moderate", variant="primary")
gr.Examples(["have a nice day!", "you are a worthless idiot", "i will find you and hurt you",
"WIN a FREE iPhone, click http://bit.ly/x now"], inp)
with gr.Column():
verdict = gr.Markdown()
scores = gr.Label(label="Scores", num_top_classes=13)
raw = gr.JSON(label="Response")
btn.click(ui, inp, [verdict, scores, raw], api_name=False)
inp.submit(ui, inp, [verdict, scores, raw], api_name=False)
gr.api(moderate, api_name="moderate")
gr.api(moderate_batch, api_name="moderate_batch")
demo.queue(default_concurrency_limit=4)
app = FastAPI(title="Moderation API")
from fastapi.middleware.cors import CORSMiddleware
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
class Req(BaseModel):
input: str | list[str] = Field(..., description="A text or a list of up to 32 texts")
@app.post("/v1/moderate")
def rest(req: Req):
texts = [req.input] if isinstance(req.input, str) else req.input
return {"model": REPO, "results": _run(texts)}
@app.get("/v1/labels")
def labels():
return {"labels": LABELS, "thresholds": TH}
app = gr.mount_gradio_app(app, demo, path="/")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)
|