infer_hf.py

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()