ncii_guard.py

"""The NCII prompt guard, in its own process.

[`hfmlsoc/ncii-guard-v02`](https://huggingface.co/hfmlsoc/ncii-guard-v02) is a 270M CPU text classifier
scoring the NCII risk of an edit prompt. It cannot live in the main process: with it loaded there, every
subsequent `@spaces.GPU` worker dies at `worker_init` with `RuntimeError: No CUDA GPUs are available` — the
fork inherits whatever CUDA driver state the classifier's torch activity left behind, and a restart does not
clear it. It cannot be a `multiprocessing.spawn` child either: spawn re-imports the parent's main module, and
on a Space that main module is `app.py` — the child would re-run the whole startup, model load included. So
the classifier runs this file as a plain subprocess — a fresh interpreter that never sees `spaces` — and
answers over stdin/stdout, one JSON object per line.
"""

from __future__ import annotations

import json
import os
import select
import subprocess
import sys
import threading

GUARD_REPO = "hfmlsoc/ncii-guard-v02"
# Prompts are short; the model's own max_position_embeddings is 32768, and cost
# grows with length (43ms at 16 tokens, 267ms at 256), so cap it.
MAX_TOKENS = 128
# A 270M model over many threads is pathologically slow — 6.1s per prompt on 30
# threads vs 41ms on 8 (OpenMP oversubscription on small tensors).
CPU_THREADS = 8

_lock = threading.Lock()
_process: subprocess.Popen | None = None


def _read(timeout: float) -> dict:
    readable, _, _ = select.select([_process.stdout], [], [], timeout)
    if not readable:
        raise TimeoutError(f"the guard did not answer within {timeout}s")
    line = _process.stdout.readline()
    if not line:
        raise EOFError("the guard process died")
    return json.loads(line)


def _spawn() -> None:
    global _process
    env = dict(os.environ)
    # Belt and braces: this child must never touch a GPU.
    env["CUDA_VISIBLE_DEVICES"] = ""
    env["OMP_NUM_THREADS"] = str(CPU_THREADS)
    _process = subprocess.Popen(
        [sys.executable, os.path.abspath(__file__)],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        text=True,
        bufsize=1,
        env=env,
    )
    # Generous: a cold cache downloads the 1.07GB checkpoint first.
    assert _read(300.0) == {"status": "ready"}


def start() -> None:
    """Launch the worker and block until its model is up.

    Called once at startup, before the main model is loaded so the fork is cheap.
    `classify` revives it if it dies.
    """
    with _lock:
        _spawn()


def classify(prompt: str, timeout: float = 60.0) -> dict:
    """`{'p_ncii': float}` for one prompt, replacing a dead or wedged worker once."""
    with _lock:
        for attempt in (0, 1):
            try:
                if _process is None or _process.poll() is not None:
                    _spawn()
                _process.stdin.write(json.dumps({"prompt": prompt}) + "\n")
                _process.stdin.flush()
                return _read(timeout)
            except Exception:
                if attempt:
                    raise
                if _process is not None and _process.poll() is None:
                    _process.kill()


def _serve() -> None:
    """The child: plain torch on CPU. The protocol keeps the real stdout to itself — everything else
    (download progress, warnings) is pushed over to stderr so it cannot corrupt a reply."""
    protocol = os.fdopen(os.dup(1), "w", buffering=1)
    os.dup2(2, 1)

    import torch
    from transformers import pipeline

    torch.set_num_threads(min(CPU_THREADS, os.cpu_count() or CPU_THREADS))
    # The tokenizer must come from this repo: it carries the confusable/bidi
    # normalizer that folds homoglyph obfuscation ("rеmove" with a Cyrillic е).
    classifier = pipeline("text-classification", model=GUARD_REPO, device="cpu")
    protocol.write(json.dumps({"status": "ready"}) + "\n")
    for line in sys.stdin:
        prompt = json.loads(line)["prompt"]
        # top_k=None returns every label, so the caller can threshold on the
        # ncii probability itself rather than on the argmax label.
        scores = classifier(prompt, truncation=True, max_length=MAX_TOKENS, top_k=None)
        p_ncii = 0.0
        for entry in scores:
            if str(entry["label"]).lower() == "ncii":
                p_ncii = float(entry["score"])
                break
        protocol.write(json.dumps({"p_ncii": p_ncii}) + "\n")


if __name__ == "__main__":
    _serve()