Deploy Threat Detection Model

import json
import logging
import math
import os

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

logger = logging.getLogger(__name__)

model = None
tokenizer = None
device = None

MAX_TOKENS = 512
MAX_CHUNKS = 32
BATCH_SIZE = 16
TEMPERATURE = 1.0


def init():
    global model, tokenizer, device
    model_dir = os.getenv("AZUREML_MODEL_DIR", "./model")
    if not os.path.isfile(os.path.join(model_dir, "config.json")):
        subdirs = [d for d in os.listdir(model_dir)
                   if os.path.isdir(os.path.join(model_dir, d))]
        if len(subdirs) == 1:
            model_dir = os.path.join(model_dir, subdirs[0])
    logger.info(f"Loading model from {model_dir}")
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForSequenceClassification.from_pretrained(model_dir)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    model.eval()
    logger.info(f"Model loaded on {device}, labels: {model.config.id2label}")


def _score(content):
    tokens = tokenizer(content, truncation=False,
                       add_special_tokens=False)["input_ids"]
    if len(tokens) <= MAX_TOKENS:
        chunks = [content]
    else:
        chunks = [
            tokenizer.decode(tokens[i : i + MAX_TOKENS], skip_special_tokens=True)
            for i in range(0, len(tokens), MAX_TOKENS)
        ]
        if len(chunks) > MAX_CHUNKS:
            logger.warning(
                f"Content produced {len(chunks)} chunks, scoring first {MAX_CHUNKS}"
            )
            chunks = chunks[:MAX_CHUNKS]

    chunks = [c for c in chunks if c]
    if not chunks:
        logger.warning("No scoreable chunks after filtering, treating as BENIGN")
        return {"label": "BENIGN", "score": 0.0}

    scores = []
    for i in range(0, len(chunks), BATCH_SIZE):
        inputs = tokenizer(
            chunks[i : i + BATCH_SIZE],
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=MAX_TOKENS,
        ).to(device)
        with torch.no_grad():
            logits = model(**inputs).logits
        probs = torch.softmax(logits / TEMPERATURE, dim=-1)
        scores.extend(probs[:, 1].tolist())

    score = max(scores)
    if len(chunks) > 1:
        score = score ** math.sqrt(len(chunks))
    label = model.config.id2label[1 if score >= 0.5 else 0]
    return {"label": label, "score": round(score, 6)}


def run(raw_data):
    items = json.loads(raw_data).get("data", [])
    return [
        _score(item["content"]) if item.get("content")
        else {"label": "BENIGN", "score": 0.0}
        for item in items
    ]