File size: 1,735 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 47 | import argparse
import json
from pathlib import Path
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
def main():
parser = argparse.ArgumentParser(description="Classify a text log file one line at a time.")
parser.add_argument("model", help="Local model directory or Hugging Face repo id")
parser.add_argument("input", help="Input text log file")
parser.add_argument("--output", default="classified.jsonl")
parser.add_argument("--threshold", type=float, default=0.5)
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()
input_path = Path(args.input)
output_path = Path(args.output)
with input_path.open("r", encoding="utf-8", errors="replace") as src, output_path.open("w", encoding="utf-8") as dst:
for line_no, raw in enumerate(src, 1):
text = raw.rstrip("\r\n")
if not text:
continue
encoded = tokenizer(text, return_tensors="pt", truncation=True, max_length=96)
with torch.inference_mode():
probs = torch.softmax(model(**encoded).logits, dim=-1)[0]
suspicious = float(probs[1])
label = "SUSPICIOUS" if suspicious >= args.threshold else "BENIGN"
dst.write(json.dumps({
"line": line_no,
"label": label,
"suspicious_probability": round(suspicious, 6),
"text": text,
}) + "\n")
print(f"Wrote {output_path}")
if __name__ == "__main__":
main()
|