File size: 1,209 Bytes
12097aa | 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 | import argparse
import json
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
def main():
parser = argparse.ArgumentParser()
parser.add_argument("model", help="Local model directory or Hugging Face repo id")
parser.add_argument("text", nargs="?", default="EventID=4625 Failed logon user=administrator source_ip=203.0.113.44 count=17")
args = parser.parse_args()
torch.set_num_threads(2)
tokenizer = AutoTokenizer.from_pretrained(
args.model,
trust_remote_code=True,
)
model = AutoModelForSequenceClassification.from_pretrained(
args.model,
trust_remote_code=True,
)
model.eval()
encoded = tokenizer(
args.text,
return_tensors="pt",
truncation=True,
max_length=96,
padding=False,
)
with torch.inference_mode():
logits = model(**encoded).logits
probs = torch.softmax(logits, dim=-1)[0]
idx = int(probs.argmax().item())
print(json.dumps({
"label": model.config.id2label[idx],
"confidence": round(float(probs[idx]), 6),
"text": args.text,
}, indent=2))
if __name__ == "__main__":
main()
|