Spaces:
Runtime error
Runtime error
File size: 4,192 Bytes
d4ec2fa | 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 | """
Loads the fine-tuned ATA-chapter and severity models ONCE at import time (not per-request
-- reloading a transformer per request would blow past the 2-second latency target), and
exposes simple predict functions for the route handlers to call.
"""
import json
import logging
import torch
import torch.nn.functional as F
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from api.config import ATA_MODEL_PATH, SEVERITY_MODEL_PATH, MAX_NARRATIVE_TOKENS
logger = logging.getLogger(__name__)
class ModelService:
def __init__(self):
self.ata_tokenizer = None
self.ata_model = None
self.ata_id2label = None
self.severity_tokenizer = None
self.severity_model = None
self.severity_id2label = None
self.loaded = False
self.load_error = None
def load(self):
"""Called once at API startup. Failure here is what triggers the 503 responses
described in the Security doc ('Model service unavailable / fails to load')."""
try:
self.ata_tokenizer = AutoTokenizer.from_pretrained(str(ATA_MODEL_PATH), local_files_only=True)
self.ata_model = AutoModelForSequenceClassification.from_pretrained(
str(ATA_MODEL_PATH), local_files_only=True, low_cpu_mem_usage=True
)
self.ata_model.eval()
with open(ATA_MODEL_PATH / "label_map.json") as f:
self.ata_id2label = {int(k): v for k, v in json.load(f)["id2label"].items()}
self.severity_tokenizer = AutoTokenizer.from_pretrained(
str(SEVERITY_MODEL_PATH), local_files_only=True
)
self.severity_model = AutoModelForSequenceClassification.from_pretrained(
str(SEVERITY_MODEL_PATH), local_files_only=True, low_cpu_mem_usage=True
)
self.severity_model.eval()
with open(SEVERITY_MODEL_PATH / "label_map.json") as f:
self.severity_id2label = {int(k): v for k, v in json.load(f)["id2label"].items()}
self.loaded = True
logger.info("Both models loaded successfully.")
except Exception as e:
self.loaded = False
self.load_error = str(e)
logger.error(f"Model loading failed: {e}")
def _predict_single(self, tokenizer, model, id2label, text: str, top_k: int = 1):
was_truncated = False
token_count = len(tokenizer.encode(text))
if token_count > MAX_NARRATIVE_TOKENS:
was_truncated = True
inputs = tokenizer(text, truncation=True, max_length=MAX_NARRATIVE_TOKENS, return_tensors="pt")
with torch.no_grad():
logits = model(**inputs).logits
probs = F.softmax(logits, dim=1)[0]
top_probs, top_ids = torch.topk(probs, k=min(top_k, len(id2label)))
results = [
{"label": id2label[int(idx)], "confidence": float(p)}
for p, idx in zip(top_probs.tolist(), top_ids.tolist())
]
return results, was_truncated
def predict(self, text: str):
if not self.loaded:
raise RuntimeError(self.load_error or "Models not loaded")
# top 3 for ATA so the API can expose "other possible systems" per the
# Security doc's ambiguous multi-system report handling
ata_results, ata_truncated = self._predict_single(
self.ata_tokenizer, self.ata_model, self.ata_id2label, text, top_k=3
)
severity_results, severity_truncated = self._predict_single(
self.severity_tokenizer, self.severity_model, self.severity_id2label, text, top_k=1
)
return {
"ata_chapter": ata_results[0]["label"],
"ata_confidence": ata_results[0]["confidence"],
"ata_other_possible": ata_results[1:], # next-highest scoring classes
"severity": severity_results[0]["label"],
"severity_confidence": severity_results[0]["confidence"],
"text_was_truncated": ata_truncated or severity_truncated,
}
# Singleton -- one instance shared across all requests, loaded once at startup
model_service = ModelService()
|