| 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() |
| } |
|
|
| |
| |
| 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 |
| ], |
| } |
|
|
