#!/usr/bin/env python3
"""Loader and inference for the Opus Moderation models.

Three of the four load with plain `from_pretrained`. mod-3 does not - it is a
LoRA adapter over Qwen3-1.7B that was converted from a decoder into an encoder
at training time, and inference has to rebuild that conversion exactly or the
weights mean nothing. `ModerationModel.load()` handles both paths.

    from opus_moderation import ModerationModel

    mod = ModerationModel.load("adapters/om4-large")
    neo = ModerationModel.load("adapters/om4-neo")

    mod.score(["some text"])   # -> [{"toxicity": 0.02, ...}]
    neo.flag(["some text"])    # -> [{"jailbreaking": False}]

Run it directly for a smoke test:

    python opus_moderation.py adapters/om4-large "you are a helpful assistant"
"""
import json
from pathlib import Path

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# The eight labels, in the order the models emit them. om4-neo is the
# exception: it has one output, "jailbreaking". Always read id2label off the
# loaded config rather than assuming this list.
LABELS = ["toxicity", "severe_toxicity", "obscene", "threat", "insult",
          "identity_attack", "sexual_explicit", "jailbreaking"]

# Used only when a model ships no thresholds.json. severe_toxicity sits at 0.2
# because its positives are extremely rare - at 0.5 it fires on nothing.
DEFAULT_THRESHOLDS = {c: 0.5 for c in LABELS}
DEFAULT_THRESHOLDS["severe_toxicity"] = 0.2


def _build_qwen_encoder(path, arch):
    """Rebuild the decoder-to-encoder conversion mod-3 was trained under.

    Two separate changes, and the second one is not optional. Swapping the
    causal mask builder is not enough on its own: for an unpadded batch the
    mask builders return None, and the SDPA path then falls back to
    `module.is_causal` - which Qwen3Attention hardcodes to True. Skip that step
    and the model stays fully causal while appearing to load fine, quietly
    scoring every input wrong.
    """
    import transformers.models.qwen3.modeling_qwen3 as qm
    from peft import PeftModel
    from transformers.masking_utils import (
        create_bidirectional_mask, create_bidirectional_sliding_window_mask)
    from transformers.modeling_outputs import SequenceClassifierOutputWithPast
    from transformers.models.qwen3.modeling_qwen3 import Qwen3ForSequenceClassification

    if arch.get("bidirectional"):
        qm.create_causal_mask = create_bidirectional_mask
        qm.create_sliding_window_causal_mask = create_bidirectional_sliding_window_mask

    class MeanPool(Qwen3ForSequenceClassification):
        def forward(self, input_ids=None, attention_mask=None, **kw):
            for k in ("labels", "num_items_in_batch", "task_ids"):
                kw.pop(k, None)
            out = self.model(input_ids=input_ids, attention_mask=attention_mask)
            h = out.last_hidden_state
            if attention_mask is None:
                pooled = h.mean(1)
            else:
                m = attention_mask.unsqueeze(-1).to(h.dtype)
                pooled = (h * m).sum(1) / m.sum(1).clamp(min=1e-6)
            return SequenceClassifierOutputWithPast(loss=None, logits=self.score(pooled))

    cls = MeanPool if arch.get("pool") == "mean" else Qwen3ForSequenceClassification
    base = cls.from_pretrained(
        arch["base_model"], num_labels=len(LABELS),
        problem_type="multi_label_classification", attn_implementation="sdpa",
        dtype=torch.bfloat16)

    if arch.get("bidirectional"):
        flipped = 0
        for m in base.modules():
            if getattr(m, "is_causal", False):
                m.is_causal = False
                flipped += 1
        if flipped == 0:
            raise RuntimeError(
                "no attention module exposed is_causal - transformers internals "
                "have moved and this model would silently score as causal.")

    keep = arch.get("kept_layers")
    cfg = base.config.get_text_config()
    if keep and keep < cfg.num_hidden_layers:
        base.model.layers = base.model.layers[:keep]
        cfg.num_hidden_layers = keep
        if getattr(cfg, "layer_types", None):
            cfg.layer_types = cfg.layer_types[:keep]

    return PeftModel.from_pretrained(base, str(path))


class ModerationModel:
    def __init__(self, model, tok, labels, thresholds, name):
        self.model, self.tok = model, tok
        self.labels, self.thresholds, self.name = labels, thresholds, name
        self.model.eval()

    @classmethod
    def load(cls, path, device=None):
        path = Path(path)
        device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        arch_file = path / "opus_arch.json"
        arch = json.loads(arch_file.read_text()) if arch_file.exists() else {}

        # A LoRA adapter directory has no model weights of its own. That, not
        # the name, is what decides which loading path to take.
        is_adapter = (path / "adapter_config.json").exists()
        if is_adapter or arch.get("base_model", "").startswith("Qwen"):
            model = _build_qwen_encoder(path, arch)
            tok_src = arch.get("base_model", str(path))
        else:
            model = AutoModelForSequenceClassification.from_pretrained(path)
            tok_src = str(path)

        tok = AutoTokenizer.from_pretrained(tok_src)
        if tok.pad_token is None:
            tok.pad_token = tok.eos_token
        model.config.pad_token_id = tok.pad_token_id
        model.to(device)

        id2label = getattr(model.config, "id2label", None) or {}
        labels = ([id2label[i] for i in sorted(id2label)]
                  if len(id2label) == model.config.num_labels else LABELS)

        thr_file = path / "thresholds.json"
        thresholds = json.loads(thr_file.read_text()) if thr_file.exists() else {}
        thresholds = {lab: thresholds.get(lab, DEFAULT_THRESHOLDS.get(lab, 0.5))
                      for lab in labels}
        return cls(model, tok, labels, thresholds, path.name)

    @torch.no_grad()
    def score(self, texts, max_len=512, batch=32):
        """Raw probabilities per label. Sorting by length first cuts padding -
        it was worth ~5x on batch inference during evaluation."""
        if isinstance(texts, str):
            texts = [texts]
        order = sorted(range(len(texts)), key=lambda i: -len(texts[i]))
        out = [None] * len(texts)
        device = next(self.model.parameters()).device
        for i in range(0, len(order), batch):
            idx = order[i:i + batch]
            enc = self.tok([texts[j] for j in idx], truncation=True,
                           max_length=max_len, padding=True,
                           return_tensors="pt").to(device)
            probs = torch.sigmoid(self.model(**enc).logits.float()).cpu()
            for row, j in zip(probs, idx):
                out[j] = {lab: float(row[k]) for k, lab in enumerate(self.labels)}
        return out

    def flag(self, texts, **kw):
        """Booleans at this model's deployment thresholds."""
        return [{lab: p[lab] >= self.thresholds[lab] for lab in self.labels}
                for p in self.score(texts, **kw)]


if __name__ == "__main__":
    import sys

    if len(sys.argv) < 3:
        print(__doc__)
        raise SystemExit(1)
    m = ModerationModel.load(sys.argv[1])
    print(f"loaded {m.name}: {len(m.labels)} labels {m.labels}")
    print(f"thresholds: {m.thresholds}\n")
    for text, probs, flags in zip(sys.argv[2:], m.score(sys.argv[2:]),
                                  m.flag(sys.argv[2:])):
        hits = [k for k, v in flags.items() if v]
        print(f"{text[:70]!r}")
        print(f"  flagged: {hits or 'nothing'}")
        print("  " + "  ".join(f"{k}={v:.3f}" for k, v in probs.items()))