Views
No views yet
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4tokenizer = AutoTokenizer.from_pretrained("cyber-threat-detector-v1")
5model = AutoModelForSequenceClassification.from_pretrained("cyber-threat-detector-v1")
6
7log_line = "SELECT * FROM users WHERE id = '1' OR '1'='1';"
8inputs = tokenizer(log_line, return_tensors="pt")
9outputs = model(**inputs)
10prediction = torch.argmax(outputs.logits)
11print("Threat Detected" if prediction == 1 else "Safe")