File size: 2,716 Bytes
9e9a98f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | import json
import os
from pathlib import Path
import torch
from transformers import (
AutoModelForSequenceClassification
)
class EndpointHandler:
def __init__(self, path=""):
folder = Path(path)
self.device = (
"cuda" if torch.cuda.is_available()
else "cpu"
)
self.model = (
AutoModelForSequenceClassification
.from_pretrained(
str(folder),
trust_remote_code=True,
local_files_only=True,
)
)
self.model.to(self.device).eval()
self.thresholds = {}
for name in ("dense", "moe"):
data = json.loads(
(folder / name / "metrics.json")
.read_text()
)
self.thresholds[name] = {
key: float(value["threshold"])
for key, value
in data["thresholds"].items()
}
# Preload gated encoder and fail startup if
# required credentials or model access are missing.
self.model._get_encoder()
@torch.inference_mode()
def __call__(self, data):
texts = data.get("inputs")
params = data.get("parameters") or {}
if isinstance(texts, str):
texts = [texts]
if not isinstance(texts, list):
raise ValueError(
"inputs must be a string or list of strings"
)
head = params.get("head", "moe")
policy = params.get(
"threshold_id", "T0_fixed_0_5"
)
if head not in ("dense", "moe"):
raise ValueError("Unsupported head")
if policy not in self.thresholds[head]:
raise ValueError("Unsupported threshold")
output = self.model(
texts=texts,
head=head,
)
scores = torch.softmax(
output.logits, dim=-1
)[:, 1].detach().cpu().tolist()
threshold = self.thresholds[head][policy]
return {
"model": "selimk92/AstraeaGuard-V2",
"head": head,
"threshold_id": policy,
"threshold": threshold,
"status": "research_only",
"blocking_enabled": False,
"results": [
{
"label": (
"ATTACK"
if score >= threshold
else "BENIGN"
),
"attack_score": float(score),
"flagged": bool(
score >= threshold
),
}
for score in scores
],
}
|