The primary goal of this model is to act as the "brain" for an autonomous security scanning agent, enabling it to understand context and make strategic decisions.
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4# Replace 'kangali/durs-llm-web-scanner' with your repo name if different
5model_name = "kangali/durs-llm-web-scanner"
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7model = AutoModelForSequenceClassification.from_pretrained(model_name)
8
9# Example inputs
10text_inputs = [
11 "<script>alert(1)</script>", # XSS Payload
12 "user_id", # IDOR Context
13 "A probe string is reflected inside an HTML tag...", # Scanner Logic
14 "This is a normal comment." # Benign
15]
16
17# Prediction
18inputs = tokenizer(text_inputs, return_tensors="pt", padding=True, truncation=True)
19with torch.no_grad():
20 logits = model(**inputs).logits
21
22predicted_class_ids = torch.argmax(logits, dim=1)
23
24# To get the string labels, you need the label_encoder.joblib from the repository
25# or you can map them manually from the model's config.json
26for i, text in enumerate(text_inputs):
27 predicted_id = predicted_class_ids[i].item()
28 label = model.config.id2label[predicted_id]
29 print(f"Input: '{text[:50]}...' -> Predicted Label: {label}")
30
This model was trained on a custom-built
master_training_dataset.csv dataset, which contains
over 2700 samples extracted and synthesized from the codebase of
Dursgo, an open-source web security scanner.
This model is a
distilbert-base-uncased that has been fine-tuned for 50 epochs using the
Trainer from the Hugging Face Transformers library. The complete workflow for creating the dataset and retraining this model is available in the project's GitHub repository:
Tunning-AI (Repo Private - To Be Continue to Open).