| """Standalone inference for the DistilBERT email-triage model (see module docstring).""" |
| INPUT_FORMAT = ("Subject: ...\nFrom: ***@domain\nReply-To: ...\nBody: ...\nURLs: ... " |
| "(§4 normalize_row + PII redaction; truncated at max_length)") |
| import json, os |
| import numpy as np |
| import torch |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification |
|
|
| MODEL_DIR = os.environ.get('JEV_MODEL_DIR', os.path.dirname(os.path.abspath(__file__))) |
| with open(os.path.join(MODEL_DIR, 'label_map.json')) as _f: |
| _LM = json.load(_f) |
| LABELS = _LM['labels'] |
| ACTION = {'legitimate': 'allow', 'spam': 'filter_to_spam', |
| 'phishing': 'quarantine_or_review', 'malware': 'quarantine_or_review'} |
| _THRESHOLD = float(_LM.get('review_threshold', 0.55)) |
| _MAX = int(_LM.get('max_length', 512)) |
| _tok = AutoTokenizer.from_pretrained(MODEL_DIR) |
| _model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR).eval() |
|
|
| @torch.no_grad() |
| def predict(texts, threshold=_THRESHOLD, batch_size=16): |
| outs, all_p = [], [] |
| for i in range(0, len(texts), batch_size): |
| t = _tok(list(texts[i:i + batch_size]), truncation=True, padding=True, |
| max_length=_MAX, return_tensors='pt') |
| all_p.append(torch.softmax(_model(**t).logits, -1).cpu().numpy()) |
| import numpy as _np |
| P = _np.concatenate(all_p) if all_p else _np.zeros((0, len(LABELS))) |
| for t, p in zip(texts, P): |
| scores = {l: float(p[LABELS.index(l)]) for l in LABELS} |
| top = LABELS[int(p.argmax())] |
| if not str(t).strip() or float(p.max()) < threshold: |
| outs.append({'label': 'review', 'scores': scores, 'action': 'review'}) |
| else: |
| outs.append({'label': top, 'scores': scores, 'action': ACTION[top]}) |
| return outs |
|
|
| if __name__ == '__main__': |
| demo = ['Subject: Meeting moved to 2 PM\nFrom: \nReply-To: \nBody: See the calendar invite.\nURLs: ', ''] |
| print(json.dumps(predict(demo), indent=2)) |
|
|
