main.py

import os, modal, torch
from fastapi import FastAPI, Body

app = modal.App("OpenCensor")

# deps
image = (
    modal.Image.debian_slim()
    .pip_install(
        "fastapi>=0.110",
        "transformers>=4.43,<4.47",
        "numpy<2",
        "torch==2.2.2",
    )
)

# env config
KMODEL_ID = os.getenv("MODEL_ID", "LikoKIko/OpenCensor-Hebrew")
KDEFAULT_THRESH = float(os.getenv("THRESH", "0.5"))
KDEFAULT_LEN = int(os.getenv("MAX_LENGTH", "256"))

# warmup: load tokenizer + model now so first call is fast
# runs during image build
def _warm():
    from transformers import AutoTokenizer, AutoModelForSequenceClassification
    AutoTokenizer.from_pretrained(KMODEL_ID)
    AutoModelForSequenceClassification.from_pretrained(KMODEL_ID, num_labels=1)
image = image.run_function(_warm)

# overview:
# /predict: check one text in a single request
# example: "wsup my g"

# /batch: check multiple texts in a single request
# example: ["wsup my g", "how are you mate?", ...]

@app.function(image=image, gpu="T4", timeout=600)  # remove gpu="T4" for CPU-only
@modal.asgi_app()
def api():
    app = FastAPI(title="OpenCensor")

    torch.set_grad_enabled(False)
    from transformers import AutoTokenizer, AutoModelForSequenceClassification

    tok = AutoTokenizer.from_pretrained(KMODEL_ID)
    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    mdl = AutoModelForSequenceClassification.from_pretrained(KMODEL_ID, num_labels=1).to(dev).eval()

    def _clean(s):
        return " ".join(str(s or "").split())  # removes over spacing

    @torch.inference_mode()
    def _score(txt, thr=KDEFAULT_THRESH, mx=KDEFAULT_LEN):
        t = _clean(txt)
        if not t: return {"prob": 0.0, "label": 0, "note": "empty"}
        b = tok(t, return_tensors="pt", truncation=True, padding=True, max_length=int(mx)).to(dev)
        p = torch.sigmoid(mdl(**b).logits).item()
        return {"prob": float(p), "label": 1 if p >= float(thr) else 0}

    @app.post("/predict")  # single sample
    async def predict(body: dict = Body(...)):
        t = body.get("text", "")
        thr = body.get("threshold", KDEFAULT_THRESH)
        mx = body.get("max_length", KDEFAULT_LEN)
        return _score(t, thr, mx)

    @app.post("/batch")  # batch: list of strings
    async def batch(body: dict = Body(...)):
        xs = body.get("texts") or []
        thr = body.get("threshold", KDEFAULT_THRESH)
        mx = body.get("max_length", KDEFAULT_LEN)
        return [_score(x, thr, mx) for x in xs]

    return app