Deploy Threat Detection Model

import json
import logging
import os

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

logger = logging.getLogger(__name__)

model = None
tokenizer = None
device = None

MAX_TOKENS = 512
TEMPERATURE = 1.0


def init():
    global model, tokenizer, device

    model_dir = os.getenv("AZUREML_MODEL_DIR", "./model")

    # Azure ML may nest model files in a subdirectory. Walk into it
    # if config.json isn't at the top level.
    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 run(raw_data):
    data = json.loads(raw_data)
    items = data.get("data", [])
    results = []

    for item in items:
        content = item.get("content", "")
        if not content:
            results.append({"label": "BENIGN", "score": 0.0})
            continue

        inputs = tokenizer(
            content,
            return_tensors="pt",
            truncation=True,
            max_length=MAX_TOKENS,
            padding=True,
        ).to(device)

        with torch.no_grad():
            logits = model(**inputs).logits

        probs = torch.softmax(logits / TEMPERATURE, dim=-1)
        malicious_score = probs[0, 1].item()
        label = model.config.id2label[1 if malicious_score >= 0.5 else 0]

        results.append({
            "label": label,
            "score": round(malicious_score, 6),
        })

    return results