upload model fit for web attack payload classfication/ and model based on codebert-base/ dataset used opensource

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model.onnx filter=lfs diff=lfs merge=lfs -text
+model_quantized.onnx filter=lfs diff=lfs merge=lfs -text
+*.ft filter=lfs diff=lfs merge=lfs -text

README.md

@@ -5,4 +5,361 @@ language:
 base_model:
 - microsoft/codebert-base
 pipeline_tag: text-classification
----
+---
+# Web Attack Detection Model
+
+A CodeBERT-based deep learning model for detecting malicious web requests and payloads. This model can identify SQL injection, XSS, path traversal, command injection, and other common web attack patterns.
+
+## Model Description
+
+This model is fine-tuned from [microsoft/codebert-base](https://huggingface.co/microsoft/codebert-base) for binary classification of web requests as either **benign** or **malicious**.
+
+### Model Architecture
+
+- **Base Model**: CodeBERT (RoBERTa-base architecture)
+- **Task**: Binary Text Classification
+- **Parameters**: 124.6M
+- **Max Sequence Length**: 256 tokens
+
+### Performance Metrics
+
+| Metric | Training Set | Test Set (125K) | 2000-Sample Test |
+|--------|-------------|-----------------|------------------|
+| **Accuracy** | 99.30% | 99.38% | 99.60% |
+| **Precision** | - | 99.47% | 99.80% |
+| **Recall** | - | 99.21% | 99.40% |
+| **F1 Score** | - | 99.34% | 99.60% |
+
+### Confusion Matrix (Test Set)
+
+|  | Predicted Benign | Predicted Malicious |
+|--|------------------|---------------------|
+| **Actual Benign** | 65,914 | 312 |
+| **Actual Malicious** | 464 | 58,491 |
+
+## Training Details
+
+### Dataset
+
+- **Total Samples**: 625,904
+- **Training Samples**: 500,722 (80%)
+- **Test Samples**: 125,181 (20%)
+- **Class Distribution**: Balanced (47% malicious, 53% benign)
+- **Sampling Strategy**: Balanced sampling with WeightedRandomSampler
+
+### Training Configuration
+
+| Parameter | Value |
+|-----------|-------|
+| Epochs | 3 |
+| Batch Size | 8 |
+| Gradient Accumulation Steps | 4 |
+| Effective Batch Size | 32 |
+| Learning Rate | 2e-5 |
+| Warmup Steps | 500 |
+| Weight Decay | 0.01 |
+| Max Sequence Length | 256 |
+| Optimizer | AdamW |
+
+### Training Progress
+
+| Epoch | Train Loss | Train Acc | Test Loss | Test Acc | F1 Score |
+|-------|------------|-----------|-----------|----------|----------|
+| 1 | 0.0289 | 98.84% | 0.0192 | 99.09% | 0.9904 |
+| 2 | 0.0201 | 99.24% | 0.0169 | 99.08% | 0.9903 |
+| 3 | 0.0175 | 99.30% | 0.0274 | 99.38% | 0.9934 |
+
+### Hardware
+
+- **GPU**: NVIDIA Tesla T4 (16GB)
+- **Training Time**: ~24 hours
+
+## Model Files
+
+| File | Size | Description |
+|------|------|-------------|
+| `best_model.pt` | 1.4 GB | PyTorch checkpoint (full precision) |
+| `model.onnx` | 476 MB | ONNX model (full precision) |
+| `model_quantized.onnx` | 120 MB | ONNX model (INT8 quantized) |
+
+## Usage
+
+### Quick Start with ONNX Runtime
+
+```python
+import numpy as np
+import onnxruntime as ort
+from transformers import RobertaTokenizer
+
+# Load tokenizer and model
+tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
+session = ort.InferenceSession("model_quantized.onnx", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+
+# Predict
+def predict(payload: str) -> dict:
+    inputs = tokenizer(
+        payload,
+        max_length=256,
+        padding='max_length',
+        truncation=True,
+        return_tensors='np'
+    )
+    
+    outputs = session.run(
+        None,
+        {
+            'input_ids': inputs['input_ids'].astype(np.int64),
+            'attention_mask': inputs['attention_mask'].astype(np.int64)
+        }
+    )
+    
+    probs = outputs[0][0]
+    pred_idx = np.argmax(probs)
+    
+    return {
+        "prediction": "malicious" if pred_idx == 1 else "benign",
+        "confidence": float(probs[pred_idx]),
+        "probabilities": {
+            "benign": float(probs[0]),
+            "malicious": float(probs[1])
+        }
+    }
+
+# Example usage
+result = predict("SELECT * FROM users WHERE id=1 OR 1=1--")
+print(result)
+# {'prediction': 'malicious', 'confidence': 0.9355, 'probabilities': {'benign': 0.0645, 'malicious': 0.9355}}
+```
+
+### Using PyTorch
+
+```python
+import torch
+import torch.nn as nn
+from transformers import RobertaTokenizer, RobertaModel
+
+class CodeBERTClassifier(nn.Module):
+    def __init__(self, model_path="microsoft/codebert-base", num_labels=2, dropout=0.1):
+        super().__init__()
+        self.codebert = RobertaModel.from_pretrained(model_path)
+        self.dropout = nn.Dropout(dropout)
+        self.classifier = nn.Linear(self.codebert.config.hidden_size, num_labels)
+    
+    def forward(self, input_ids, attention_mask):
+        outputs = self.codebert(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs.pooler_output
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+        return logits
+
+# Load model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = CodeBERTClassifier()
+model.load_state_dict(torch.load("best_model.pt", map_location=device))
+model.eval()
+model.to(device)
+
+# Load tokenizer
+tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
+
+# Predict
+def predict(payload: str) -> dict:
+    inputs = tokenizer(
+        payload,
+        max_length=256,
+        padding='max_length',
+        truncation=True,
+        return_tensors='pt'
+    ).to(device)
+    
+    with torch.no_grad():
+        logits = model(inputs['input_ids'], inputs['attention_mask'])
+        probs = torch.softmax(logits, dim=-1)[0]
+    
+    pred_idx = torch.argmax(probs).item()
+    
+    return {
+        "prediction": "malicious" if pred_idx == 1 else "benign",
+        "confidence": probs[pred_idx].item()
+    }
+
+# Example
+result = predict("<script>alert('xss')</script>")
+print(result)
+# {'prediction': 'malicious', 'confidence': 0.9998}
+```
+
+## FastAPI Server
+
+### Installation
+
+```bash
+pip install onnxruntime-gpu transformers fastapi uvicorn pydantic numpy
+```
+
+### Start Server
+
+```bash
+# GPU mode (recommended)
+python server_onnx.py --device gpu --quantized --port 8000
+
+# CPU mode
+python server_onnx.py --device cpu --quantized --port 8000
+```
+
+### API Endpoints
+
+#### Health Check
+```bash
+curl http://localhost:8000/health
+```
+
+#### Single Prediction
+```bash
+curl -X POST http://localhost:8000/predict \
+  -H "Content-Type: application/json" \
+  -d '{"payload": "SELECT * FROM users WHERE id=1 OR 1=1--"}'
+```
+
+Response:
+```json
+{
+  "payload": "SELECT * FROM users WHERE id=1 OR 1=1--",
+  "prediction": "malicious",
+  "confidence": 0.9355,
+  "probabilities": {"benign": 0.0645, "malicious": 0.9355},
+  "inference_time_ms": 15.23
+}
+```
+
+#### Batch Prediction
+```bash
+curl -X POST http://localhost:8000/batch_predict \
+  -H "Content-Type: application/json" \
+  -d '{"payloads": ["<script>alert(1)</script>", "GET /api/users HTTP/1.1"]}'
+```
+
+## Docker Deployment
+
+### GPU Version
+
+```dockerfile
+# Dockerfile
+FROM nvidia/cuda:11.8-cudnn8-runtime-ubuntu22.04
+
+RUN apt-get update && apt-get install -y python3 python3-pip
+RUN pip3 install onnxruntime-gpu transformers fastapi uvicorn pydantic numpy
+
+WORKDIR /app
+COPY model_quantized.onnx ./models/
+COPY server_onnx.py .
+
+EXPOSE 8000
+CMD ["python3", "server_onnx.py", "--device", "gpu", "--quantized"]
+```
+
+### CPU Version
+
+```dockerfile
+# Dockerfile.cpu
+FROM python:3.10-slim
+
+RUN pip install onnxruntime transformers fastapi uvicorn pydantic numpy
+
+WORKDIR /app
+COPY model_quantized.onnx ./models/
+COPY server_onnx.py .
+
+EXPOSE 8000
+CMD ["python", "server_onnx.py", "--device", "cpu", "--quantized"]
+```
+
+### Docker Compose
+
+```yaml
+version: '3.8'
+services:
+  web-attack-detector:
+    build: .
+    ports:
+      - "8000:8000"
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]
+```
+
+## Attack Types Detected
+
+This model can detect various web attack patterns including:
+
+| Attack Type | Example |
+|-------------|---------|
+| **SQL Injection** | `' OR '1'='1' --` |
+| **Cross-Site Scripting (XSS)** | `<script>alert(document.cookie)</script>` |
+| **Path Traversal** | `../../etc/passwd` |
+| **Command Injection** | `; cat /etc/passwd` |
+| **LDAP Injection** | `*)(uid=*))(|(uid=*` |
+| **XML Injection** | `<?xml version="1.0"?><!DOCTYPE foo>` |
+| **Server-Side Template Injection** | `{{7*7}}` |
+
+## Limitations
+
+- The model is trained on specific attack patterns and may not detect novel or obfuscated attacks
+- Maximum input length is 256 tokens; longer payloads will be truncated
+- The model may have false positives on legitimate requests that resemble attack patterns
+- Performance may vary on different types of web applications
+
+## Ethical Considerations
+
+This model is intended for **defensive security purposes only**, including:
+- Web Application Firewalls (WAF)
+- Intrusion Detection Systems (IDS)
+- Security monitoring and alerting
+- Penetration testing and security assessments
+
+**Do not use this model for malicious purposes.**
+
+## License
+
+This model is released under the MIT License.
+
+## Citation
+
+If you use this model in your research or application, please cite:
+
+```bibtex
+@misc{web-attack-detection-codebert,
+  author = {Your Name},
+  title = {Web Attack Detection Model based on CodeBERT},
+  year = {2024},
+  publisher = {Hugging Face},
+  howpublished = {\url{https://huggingface.co/your-username/web-attack-detection}},
+  note = {Fine-tuned CodeBERT model for detecting malicious web requests}
+}
+
+@article{feng2020codebert,
+  title = {CodeBERT: A Pre-Trained Model for Programming and Natural Languages},
+  author = {Feng, Zhangyin and Guo, Daya and Tang, Duyu and Duan, Nan and Feng, Xiaocheng and Gong, Ming and Shou, Linjun and Qin, Bing and Liu, Ting and Jiang, Daxin and Zhou, Ming},
+  journal = {Findings of the Association for Computational Linguistics: EMNLP 2020},
+  year = {2020},
+  pages = {1536--1547},
+  doi = {10.18653/v1/2020.findings-emnlp.139}
+}
+
+@article{liu2019roberta,
+  title = {RoBERTa: A Robustly Optimized BERT Pretraining Approach},
+  author = {Liu, Yinhan and Ott, Myle and Goyal, Naman and Du, Jingfei and Joshi, Mandar and Chen, Danqi and Levy, Omer and Lewis, Mike and Zettlemoyer, Luke and Stoyanov, Veselin},
+  journal = {arXiv preprint arXiv:1907.11692},
+  year = {2019}
+}
+```
+
+## Acknowledgments
+
+- [Microsoft CodeBERT](https://github.com/microsoft/CodeBERT) for the pre-trained model
+- [Hugging Face Transformers](https://huggingface.co/transformers/) for the model framework
+- [ONNX Runtime](https://onnxruntime.ai/) for efficient inference

best_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c9bce29b63361bf2b8c0554f4a51b7161c4cccad0ac784f06d4e0a435116f3d
+size 1495970402

export_onnx_quantized.py

@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+"""
+Export trained CodeBERT model to ONNX format with optional quantization.
+Supports both CPU and GPU inference.
+"""
+
+import os
+import sys
+import torch
+import torch.nn as nn
+from transformers import RobertaTokenizer, RobertaModel
+import json
+
+# Paths
+MODEL_PATH = "/c1/new-models/best_model.pt"
+CODEBERT_PATH = "/c1/huggingface/codebert-base"
+OUTPUT_DIR = "/c1/new-models"
+ONNX_PATH = os.path.join(OUTPUT_DIR, "model.onnx")
+ONNX_QUANTIZED_PATH = os.path.join(OUTPUT_DIR, "model_quantized.onnx")
+
+
+class CodeBERTClassifier(nn.Module):
+    """CodeBERT-based classifier for web attack detection - matches training script."""
+    
+    def __init__(self, model_path, num_labels=2, dropout=0.1):
+        super(CodeBERTClassifier, self).__init__()
+        self.codebert = RobertaModel.from_pretrained(model_path)
+        self.dropout = nn.Dropout(dropout)
+        self.classifier = nn.Linear(self.codebert.config.hidden_size, num_labels)
+    
+    def forward(self, input_ids, attention_mask):
+        outputs = self.codebert(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs.pooler_output
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+        return logits
+
+
+class ONNXCodeBERTClassifier(nn.Module):
+    """Wrapper for ONNX export with softmax output."""
+    
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        self.model.dropout.p = 0  # Disable dropout for inference
+    
+    def forward(self, input_ids, attention_mask):
+        outputs = self.model.codebert(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs.pooler_output
+        logits = self.model.classifier(pooled_output)
+        probabilities = torch.softmax(logits, dim=-1)
+        return probabilities
+
+
+def export_to_onnx():
+    """Export model to ONNX format."""
+    print("=" * 80)
+    print("ONNX Model Export")
+    print("=" * 80)
+    
+    # Device - use CPU for export to avoid CUDA issues
+    device = torch.device("cpu")
+    print(f"Export Device: {device}")
+    
+    # Load tokenizer
+    print("\n1. Loading tokenizer...")
+    tokenizer = RobertaTokenizer.from_pretrained(CODEBERT_PATH)
+    print(f"   Tokenizer loaded: {type(tokenizer).__name__}")
+    
+    # Create model with same architecture as training
+    print("\n2. Loading model...")
+    model = CodeBERTClassifier(CODEBERT_PATH)
+    
+    # Load trained weights
+    checkpoint = torch.load(MODEL_PATH, map_location=device)
+    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+        model.load_state_dict(checkpoint['model_state_dict'])
+    else:
+        model.load_state_dict(checkpoint)
+    
+    model.eval()
+    model.to(device)
+    print(f"   Model loaded from: {MODEL_PATH}")
+    
+    # Wrap for ONNX export
+    onnx_model = ONNXCodeBERTClassifier(model)
+    onnx_model.eval()
+    onnx_model.to(device)
+    
+    # Create dummy input
+    print("\n3. Creating dummy input...")
+    max_length = 256
+    dummy_text = "SELECT * FROM users WHERE id=1"
+    inputs = tokenizer(
+        dummy_text,
+        max_length=max_length,
+        padding='max_length',
+        truncation=True,
+        return_tensors='pt'
+    )
+    
+    dummy_input_ids = inputs['input_ids'].to(device)
+    dummy_attention_mask = inputs['attention_mask'].to(device)
+    print(f"   Input shape: {dummy_input_ids.shape}")
+    
+    # Test forward pass first
+    print("\n4. Testing forward pass...")
+    with torch.no_grad():
+        test_output = onnx_model(dummy_input_ids, dummy_attention_mask)
+        print(f"   Output shape: {test_output.shape}")
+        print(f"   Output sample: {test_output[0].numpy()}")
+    
+    # Export to ONNX
+    print("\n5. Exporting to ONNX...")
+    torch.onnx.export(
+        onnx_model,
+        (dummy_input_ids, dummy_attention_mask),
+        ONNX_PATH,
+        export_params=True,
+        opset_version=14,
+        do_constant_folding=True,
+        input_names=['input_ids', 'attention_mask'],
+        output_names=['probabilities'],
+        dynamic_axes={
+            'input_ids': {0: 'batch_size', 1: 'sequence_length'},
+            'attention_mask': {0: 'batch_size', 1: 'sequence_length'},
+            'probabilities': {0: 'batch_size'}
+        }
+    )
+    
+    onnx_size = os.path.getsize(ONNX_PATH) / (1024 * 1024)
+    print(f"   ONNX model saved: {ONNX_PATH}")
+    print(f"   Size: {onnx_size:.2f} MB")
+    
+    # Quantize model
+    print("\n6. Quantizing model (dynamic quantization)...")
+    try:
+        from onnxruntime.quantization import quantize_dynamic, QuantType
+        
+        quantize_dynamic(
+            model_input=ONNX_PATH,
+            model_output=ONNX_QUANTIZED_PATH,
+            weight_type=QuantType.QUInt8,
+            optimize_model=True
+        )
+        
+        quantized_size = os.path.getsize(ONNX_QUANTIZED_PATH) / (1024 * 1024)
+        print(f"   Quantized model saved: {ONNX_QUANTIZED_PATH}")
+        print(f"   Size: {quantized_size:.2f} MB")
+        print(f"   Compression ratio: {onnx_size / quantized_size:.2f}x")
+        
+    except Exception as e:
+        print(f"   Warning: Quantization failed: {e}")
+        print("   Using non-quantized model.")
+        import shutil
+        shutil.copy(ONNX_PATH, ONNX_QUANTIZED_PATH)
+    
+    # Verify ONNX model
+    print("\n7. Verifying ONNX model...")
+    try:
+        import onnx
+        onnx_check = onnx.load(ONNX_PATH)
+        onnx.checker.check_model(onnx_check)
+        print("   ONNX model verification: PASSED")
+    except Exception as e:
+        print(f"   Warning: ONNX verification failed: {e}")
+    
+    # Test inference with ONNX Runtime
+    print("\n8. Testing ONNX Runtime inference...")
+    try:
+        import onnxruntime as ort
+        import numpy as np
+        
+        # Try GPU first, fallback to CPU
+        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+        available_providers = ort.get_available_providers()
+        use_providers = [p for p in providers if p in available_providers]
+        
+        session = ort.InferenceSession(ONNX_PATH, providers=use_providers)
+        actual_provider = session.get_providers()[0]
+        print(f"   Using provider: {actual_provider}")
+        
+        # Test inference
+        test_texts = [
+            "SELECT * FROM users WHERE id=1 OR 1=1",  # SQL injection
+            "GET /index.html HTTP/1.1",  # Normal request
+            "<script>alert('xss')</script>",  # XSS
+            "Mozilla/5.0 (Windows NT 10.0; Win64)",  # Normal UA
+        ]
+        
+        print("\n   Test predictions:")
+        for text in test_texts:
+            inputs = tokenizer(
+                text,
+                max_length=max_length,
+                padding='max_length',
+                truncation=True,
+                return_tensors='np'
+            )
+            
+            outputs = session.run(
+                None,
+                {
+                    'input_ids': inputs['input_ids'].astype(np.int64),
+                    'attention_mask': inputs['attention_mask'].astype(np.int64)
+                }
+            )
+            
+            probs = outputs[0][0]
+            pred = np.argmax(probs)
+            label = "Malicious" if pred == 1 else "Benign"
+            conf = probs[pred] * 100
+            print(f"   - '{text[:40]:<40}' => {label:<10} ({conf:.1f}%)")
+        
+    except Exception as e:
+        print(f"   Warning: ONNX Runtime test failed: {e}")
+        import traceback
+        traceback.print_exc()
+    
+    # Save export config
+    print("\n9. Saving export configuration...")
+    export_config = {
+        "model_path": ONNX_PATH,
+        "quantized_model_path": ONNX_QUANTIZED_PATH,
+        "max_length": max_length,
+        "tokenizer_path": CODEBERT_PATH,
+        "labels": {"0": "benign", "1": "malicious"},
+        "input_names": ["input_ids", "attention_mask"],
+        "output_names": ["probabilities"]
+    }
+    
+    config_path = os.path.join(OUTPUT_DIR, "onnx_config.json")
+    with open(config_path, 'w') as f:
+        json.dump(export_config, f, indent=2)
+    print(f"   Config saved: {config_path}")
+    
+    print("\n" + "=" * 80)
+    print("Export completed!")
+    print("=" * 80)
+    print(f"ONNX Model: {ONNX_PATH} ({onnx_size:.2f} MB)")
+    if os.path.exists(ONNX_QUANTIZED_PATH):
+        qsize = os.path.getsize(ONNX_QUANTIZED_PATH) / (1024 * 1024)
+        print(f"Quantized Model: {ONNX_QUANTIZED_PATH} ({qsize:.2f} MB)")
+    print("=" * 80)
+
+
+if __name__ == "__main__":
+    export_to_onnx()

model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cbb6ba2f597dd1ec8cfd8baf6ed4df932c5fda6a20680dd093dc35c6480f712
+size 498886238

model_quantized.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ac1755d9c22be619022d59ec863c361f44646c133d7986b22e16a616fadfaba
+size 125354544

onnx_config.json

@@ -0,0 +1,17 @@
+{
+  "model_path": "/c1/new-models/model.onnx",
+  "quantized_model_path": "/c1/new-models/model_quantized.onnx",
+  "max_length": 256,
+  "tokenizer_path": "/c1/huggingface/codebert-base",
+  "labels": {
+    "0": "benign",
+    "1": "malicious"
+  },
+  "input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "output_names": [
+    "probabilities"
+  ]
+}

requirements_onnx.txt

@@ -0,0 +1,17 @@
+# ONNX Runtime Inference Requirements
+# For GPU: onnxruntime-gpu
+# For CPU: onnxruntime
+
+# Core
+onnxruntime-gpu==1.16.3
+transformers==4.35.0
+tokenizers>=0.14.0
+
+# Web framework
+fastapi==0.104.1
+uvicorn[standard]==0.24.0
+pydantic>=2.0.0
+
+# Utils
+numpy>=1.24.0
+requests>=2.31.0

requirements_onnx_cpu.txt

@@ -0,0 +1,15 @@
+# ONNX Runtime Inference Requirements (CPU only)
+
+# Core - CPU version
+onnxruntime==1.16.3
+transformers==4.35.0
+tokenizers>=0.14.0
+
+# Web framework
+fastapi==0.104.1
+uvicorn[standard]==0.24.0
+pydantic>=2.0.0
+
+# Utils
+numpy>=1.24.0
+requests>=2.31.0

server_onnx.py

@@ -0,0 +1,357 @@
+#!/usr/bin/env python3
+"""
+FastAPI server for Web Attack Detection using ONNX Runtime.
+Supports both CPU and GPU inference.
+
+Usage:
+    python server_onnx.py --host 0.0.0.0 --port 8000 --device gpu
+    python server_onnx.py --host 0.0.0.0 --port 8000 --device cpu
+    python server_onnx.py --quantized  # Use quantized model (smaller, faster)
+"""
+
+import os
+import sys
+import json
+import time
+import argparse
+import numpy as np
+from typing import List, Optional
+from contextlib import asynccontextmanager
+
+import onnxruntime as ort
+from transformers import RobertaTokenizer
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+
+# Configuration
+ONNX_MODEL_PATH = "/c1/new-models/model.onnx"
+ONNX_QUANTIZED_PATH = "/c1/new-models/model_quantized.onnx"
+TOKENIZER_PATH = "/c1/huggingface/codebert-base"
+MAX_LENGTH = 256
+
+
+class PredictRequest(BaseModel):
+    """Single prediction request."""
+    payload: str = Field(..., description="The payload/request to classify")
+
+
+class BatchPredictRequest(BaseModel):
+    """Batch prediction request."""
+    payloads: List[str] = Field(..., description="List of payloads to classify")
+
+
+class PredictResponse(BaseModel):
+    """Prediction response."""
+    payload: str
+    prediction: str  # "malicious" or "benign"
+    confidence: float
+    probabilities: dict
+    inference_time_ms: float
+
+
+class BatchPredictResponse(BaseModel):
+    """Batch prediction response."""
+    predictions: List[PredictResponse]
+    total_inference_time_ms: float
+    avg_inference_time_ms: float
+
+
+class HealthResponse(BaseModel):
+    """Health check response."""
+    status: str
+    model_loaded: bool
+    device: str
+    provider: str
+    model_path: str
+    version: str
+
+
+# Global variables
+tokenizer = None
+ort_session = None
+device_type = "cpu"
+model_path = ONNX_MODEL_PATH
+
+
+def load_model(use_gpu: bool = True, use_quantized: bool = False):
+    """Load ONNX model and tokenizer."""
+    global tokenizer, ort_session, device_type, model_path
+    
+    print("Loading model...")
+    
+    # Load tokenizer
+    print(f"  Loading tokenizer from: {TOKENIZER_PATH}")
+    tokenizer = RobertaTokenizer.from_pretrained(TOKENIZER_PATH)
+    
+    # Select model
+    model_path = ONNX_QUANTIZED_PATH if use_quantized else ONNX_MODEL_PATH
+    if not os.path.exists(model_path):
+        model_path = ONNX_MODEL_PATH
+    
+    print(f"  Loading ONNX model from: {model_path}")
+    
+    # Configure providers
+    providers = []
+    if use_gpu:
+        if 'CUDAExecutionProvider' in ort.get_available_providers():
+            providers.append('CUDAExecutionProvider')
+            device_type = "gpu"
+        else:
+            print("  Warning: CUDA not available, falling back to CPU")
+    
+    providers.append('CPUExecutionProvider')
+    if device_type != "gpu":
+        device_type = "cpu"
+    
+    # Create session
+    sess_options = ort.SessionOptions()
+    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    
+    ort_session = ort.InferenceSession(
+        model_path,
+        sess_options=sess_options,
+        providers=providers
+    )
+    
+    actual_provider = ort_session.get_providers()[0]
+    print(f"  Model loaded successfully!")
+    print(f"  Provider: {actual_provider}")
+    print(f"  Device: {device_type}")
+    
+    return ort_session
+
+
+def predict_single(payload: str) -> dict:
+    """Make prediction for a single payload."""
+    global tokenizer, ort_session
+    
+    start_time = time.time()
+    
+    # Tokenize
+    inputs = tokenizer(
+        payload,
+        max_length=MAX_LENGTH,
+        padding='max_length',
+        truncation=True,
+        return_tensors='np'
+    )
+    
+    # Run inference
+    outputs = ort_session.run(
+        None,
+        {
+            'input_ids': inputs['input_ids'].astype(np.int64),
+            'attention_mask': inputs['attention_mask'].astype(np.int64)
+        }
+    )
+    
+    # Process results
+    probs = outputs[0][0]
+    pred_idx = int(np.argmax(probs))
+    confidence = float(probs[pred_idx])
+    prediction = "malicious" if pred_idx == 1 else "benign"
+    
+    inference_time = (time.time() - start_time) * 1000
+    
+    return {
+        "payload": payload[:100] + "..." if len(payload) > 100 else payload,
+        "prediction": prediction,
+        "confidence": round(confidence, 4),
+        "probabilities": {
+            "benign": round(float(probs[0]), 4),
+            "malicious": round(float(probs[1]), 4)
+        },
+        "inference_time_ms": round(inference_time, 2)
+    }
+
+
+def predict_batch(payloads: List[str]) -> dict:
+    """Make predictions for a batch of payloads."""
+    global tokenizer, ort_session
+    
+    start_time = time.time()
+    
+    # Tokenize batch
+    inputs = tokenizer(
+        payloads,
+        max_length=MAX_LENGTH,
+        padding='max_length',
+        truncation=True,
+        return_tensors='np'
+    )
+    
+    # Run inference
+    outputs = ort_session.run(
+        None,
+        {
+            'input_ids': inputs['input_ids'].astype(np.int64),
+            'attention_mask': inputs['attention_mask'].astype(np.int64)
+        }
+    )
+    
+    total_time = (time.time() - start_time) * 1000
+    
+    # Process results
+    predictions = []
+    probs_batch = outputs[0]
+    
+    for i, (payload, probs) in enumerate(zip(payloads, probs_batch)):
+        pred_idx = int(np.argmax(probs))
+        confidence = float(probs[pred_idx])
+        prediction = "malicious" if pred_idx == 1 else "benign"
+        
+        predictions.append({
+            "payload": payload[:100] + "..." if len(payload) > 100 else payload,
+            "prediction": prediction,
+            "confidence": round(confidence, 4),
+            "probabilities": {
+                "benign": round(float(probs[0]), 4),
+                "malicious": round(float(probs[1]), 4)
+            },
+            "inference_time_ms": round(total_time / len(payloads), 2)
+        })
+    
+    return {
+        "predictions": predictions,
+        "total_inference_time_ms": round(total_time, 2),
+        "avg_inference_time_ms": round(total_time / len(payloads), 2)
+    }
+
+
+# Startup/shutdown events
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Load model on startup
+    use_gpu = getattr(app.state, 'use_gpu', True)
+    use_quantized = getattr(app.state, 'use_quantized', False)
+    load_model(use_gpu=use_gpu, use_quantized=use_quantized)
+    yield
+    # Cleanup on shutdown
+    print("Shutting down...")
+
+
+# Create FastAPI app
+app = FastAPI(
+    title="Web Attack Detection API",
+    description="CodeBERT-based web attack detection using ONNX Runtime. Supports CPU and GPU inference.",
+    version="2.0.0",
+    lifespan=lifespan
+)
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.get("/", response_model=dict)
+async def root():
+    """API root endpoint."""
+    return {
+        "name": "Web Attack Detection API",
+        "version": "2.0.0",
+        "model": "CodeBERT + ONNX Runtime",
+        "endpoints": {
+            "/predict": "POST - Single payload prediction",
+            "/batch_predict": "POST - Batch payload prediction",
+            "/health": "GET - Health check"
+        }
+    }
+
+
+@app.get("/health", response_model=HealthResponse)
+async def health():
+    """Health check endpoint."""
+    return {
+        "status": "healthy" if ort_session is not None else "unhealthy",
+        "model_loaded": ort_session is not None,
+        "device": device_type,
+        "provider": ort_session.get_providers()[0] if ort_session else "none",
+        "model_path": model_path,
+        "version": "2.0.0"
+    }
+
+
+@app.post("/predict", response_model=PredictResponse)
+async def predict(request: PredictRequest):
+    """
+    Predict if a single payload is malicious or benign.
+    
+    - **payload**: The HTTP request/payload string to analyze
+    """
+    if not ort_session:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    try:
+        result = predict_single(request.payload)
+        return result
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/batch_predict", response_model=BatchPredictResponse)
+async def batch_predict(request: BatchPredictRequest):
+    """
+    Predict if multiple payloads are malicious or benign.
+    
+    - **payloads**: List of HTTP request/payload strings to analyze
+    """
+    if not ort_session:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    if len(request.payloads) == 0:
+        raise HTTPException(status_code=400, detail="Empty payload list")
+    
+    if len(request.payloads) > 100:
+        raise HTTPException(status_code=400, detail="Maximum batch size is 100")
+    
+    try:
+        result = predict_batch(request.payloads)
+        return result
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+def main():
+    """Main entry point."""
+    parser = argparse.ArgumentParser(description="Web Attack Detection API Server")
+    parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind to")
+    parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
+    parser.add_argument("--device", type=str, default="gpu", choices=["cpu", "gpu"],
+                        help="Device to use for inference")
+    parser.add_argument("--quantized", action="store_true", 
+                        help="Use quantized model (smaller, potentially faster)")
+    parser.add_argument("--workers", type=int, default=1, help="Number of workers")
+    
+    args = parser.parse_args()
+    
+    # Store config in app state
+    app.state.use_gpu = (args.device == "gpu")
+    app.state.use_quantized = args.quantized
+    
+    print("=" * 60)
+    print("Web Attack Detection API Server")
+    print("=" * 60)
+    print(f"Host: {args.host}")
+    print(f"Port: {args.port}")
+    print(f"Device: {args.device}")
+    print(f"Quantized: {args.quantized}")
+    print("=" * 60)
+    
+    import uvicorn
+    uvicorn.run(
+        app,
+        host=args.host,
+        port=args.port,
+        workers=args.workers,
+        log_level="info"
+    )
+
+
+if __name__ == "__main__":
+    main()

test_onnx_accuracy.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+"""
+Test ONNX model accuracy with 2000 samples from the dataset.
+"""
+
+import os
+import pandas as pd
+import numpy as np
+import requests
+import time
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix, classification_report
+
+# Configuration
+API_URL = "http://localhost:8001"
+DATASET_PATH = "/c1/web-attack-detection/dataset.csv"
+NUM_SAMPLES = 2000  # 1000 malicious + 1000 benign
+BATCH_SIZE = 50
+
+
+def test_accuracy():
+    print("=" * 80)
+    print("ONNX Model Accuracy Test")
+    print("=" * 80)
+    
+    # Check API health
+    print("\n1. Checking API health...")
+    try:
+        resp = requests.get(f"{API_URL}/health", timeout=10)
+        health = resp.json()
+        print(f"   Status: {health['status']}")
+        print(f"   Device: {health['device']}")
+        print(f"   Provider: {health['provider']}")
+        print(f"   Model: {health['model_path']}")
+    except Exception as e:
+        print(f"   Error: {e}")
+        print("   Please ensure the server is running!")
+        return
+    
+    # Load dataset
+    print("\n2. Loading dataset...")
+    df = pd.read_csv(DATASET_PATH)
+    df = df.dropna(subset=['Sentence', 'Label'])
+    df['Sentence'] = df['Sentence'].astype(str)
+    df['Label'] = df['Label'].astype(int)
+    print(f"   Total samples: {len(df)}")
+    
+    # Sample data
+    print("\n3. Sampling test data...")
+    samples_per_class = NUM_SAMPLES // 2
+    
+    benign_samples = df[df['Label'] == 0].sample(n=min(samples_per_class, len(df[df['Label'] == 0])), random_state=42)
+    malicious_samples = df[df['Label'] == 1].sample(n=min(samples_per_class, len(df[df['Label'] == 1])), random_state=42)
+    
+    test_df = pd.concat([benign_samples, malicious_samples]).sample(frac=1, random_state=42).reset_index(drop=True)
+    print(f"   Test samples: {len(test_df)}")
+    print(f"   Benign: {len(test_df[test_df['Label'] == 0])}")
+    print(f"   Malicious: {len(test_df[test_df['Label'] == 1])}")
+    
+    # Run predictions
+    print("\n4. Running predictions...")
+    predictions = []
+    true_labels = []
+    total_time = 0
+    
+    for i in range(0, len(test_df), BATCH_SIZE):
+        batch = test_df.iloc[i:i+BATCH_SIZE]
+        payloads = batch['Sentence'].tolist()
+        labels = batch['Label'].tolist()
+        
+        try:
+            start = time.time()
+            resp = requests.post(
+                f"{API_URL}/batch_predict",
+                json={"payloads": payloads},
+                timeout=60
+            )
+            elapsed = time.time() - start
+            total_time += elapsed
+            
+            result = resp.json()
+            batch_preds = [1 if p['prediction'] == 'malicious' else 0 for p in result['predictions']]
+            predictions.extend(batch_preds)
+            true_labels.extend(labels)
+            
+            # Progress
+            progress = min(i + BATCH_SIZE, len(test_df))
+            print(f"   Processed: {progress}/{len(test_df)} ({100*progress/len(test_df):.1f}%)", end='\r')
+            
+        except Exception as e:
+            print(f"\n   Error at batch {i}: {e}")
+            continue
+    
+    print(f"\n   Total inference time: {total_time:.2f}s")
+    print(f"   Avg time per sample: {1000*total_time/len(predictions):.2f}ms")
+    
+    # Calculate metrics
+    print("\n5. Calculating metrics...")
+    accuracy = accuracy_score(true_labels, predictions)
+    precision = precision_score(true_labels, predictions)
+    recall = recall_score(true_labels, predictions)
+    f1 = f1_score(true_labels, predictions)
+    cm = confusion_matrix(true_labels, predictions)
+    
+    print("\n" + "=" * 80)
+    print("RESULTS")
+    print("=" * 80)
+    print(f"\nSamples tested: {len(predictions)}")
+    print(f"\nMetrics:")
+    print(f"  Accuracy:  {accuracy*100:.2f}%")
+    print(f"  Precision: {precision*100:.2f}%")
+    print(f"  Recall:    {recall*100:.2f}%")
+    print(f"  F1 Score:  {f1*100:.2f}%")
+    
+    print(f"\nConfusion Matrix:")
+    print(f"                 Predicted")
+    print(f"                 Benign  Malicious")
+    print(f"  Actual Benign    {cm[0][0]:5d}    {cm[0][1]:5d}")
+    print(f"  Actual Malicious {cm[1][0]:5d}    {cm[1][1]:5d}")
+    
+    print(f"\nDetailed Report:")
+    print(classification_report(true_labels, predictions, target_names=['Benign', 'Malicious']))
+    
+    print("=" * 80)
+    
+    # Return results
+    return {
+        'accuracy': accuracy,
+        'precision': precision,
+        'recall': recall,
+        'f1': f1,
+        'samples': len(predictions),
+        'inference_time_s': total_time
+    }
+
+
+if __name__ == "__main__":
+    test_accuracy()

tokenizer_config.json

@@ -0,0 +1,4 @@
+{
+  "model_name": "/c1/huggingface/codebert-base",
+  "max_length": 256
+}

train_new_model.py

@@ -0,0 +1,451 @@
+#!/usr/bin/env python3
+"""
+Train CodeBERT-based model for web attack detection
+Dataset: /c1/web-attack-detection/dataset.csv
+Output: /c1/new-models/
+"""
+
+import os
+import pandas as pd
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from transformers import RobertaTokenizer, RobertaModel, AdamW, get_linear_schedule_with_warmup
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support, classification_report, confusion_matrix
+from tqdm import tqdm
+import json
+import random
+from collections import Counter
+
+# Set random seeds for reproducibility
+def set_seed(seed=42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+set_seed(42)
+
+# Configuration
+class Config:
+    # Paths
+    data_path = "/c1/web-attack-detection/dataset.csv"
+    model_base_path = "/c1/huggingface/codebert-base"
+    output_dir = "/c1/new-models"
+    
+    # Training parameters
+    max_length = 256  # Reduced from 512
+    batch_size = 8    # Reduced from 32
+    gradient_accumulation_steps = 4  # Effective batch size = 8 * 4 = 32
+    epochs = 3
+    learning_rate = 2e-5
+    warmup_steps = 500
+    weight_decay = 0.01
+    
+    # Data split
+    train_size = 0.8
+    test_size = 0.2
+    
+    # Sampling strategy
+    use_sampling = True  # Enable sampling
+    sampling_strategy = "balanced"  # Options: "balanced", "oversample", "undersample", "none"
+    
+    # GPU
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Early stopping
+    early_stopping_patience = 2
+
+config = Config()
+
+print("="*80)
+print("Web Attack Detection Model Training")
+print("="*80)
+print(f"Device: {config.device}")
+print(f"Data path: {config.data_path}")
+print(f"Model base: {config.model_base_path}")
+print(f"Output dir: {config.output_dir}")
+print(f"Sampling strategy: {config.sampling_strategy}")
+print("="*80)
+
+# Create output directory
+os.makedirs(config.output_dir, exist_ok=True)
+
+# Load data
+print("\n1. Loading dataset...")
+df = pd.read_csv(config.data_path)
+print(f"Total samples: {len(df)}")
+print(f"\nLabel distribution:")
+print(df['Label'].value_counts())
+print(f"\nLabel proportions:")
+print(df['Label'].value_counts(normalize=True))
+
+# Clean data
+print("\n2. Cleaning data...")
+df = df.dropna(subset=['Sentence', 'Label'])
+df['Sentence'] = df['Sentence'].astype(str)
+df['Label'] = df['Label'].astype(int)
+print(f"Samples after cleaning: {len(df)}")
+
+# Split data
+print("\n3. Splitting data (80% train, 20% test)...")
+train_df, test_df = train_test_split(
+    df, 
+    test_size=config.test_size, 
+    random_state=42,
+    stratify=df['Label']
+)
+
+print(f"Train samples: {len(train_df)}")
+print(f"Test samples: {len(test_df)}")
+print(f"\nTrain label distribution:")
+print(train_df['Label'].value_counts())
+print(f"\nTest label distribution:")
+print(test_df['Label'].value_counts())
+
+# Apply sampling strategy
+def apply_sampling(df, strategy="balanced"):
+    """Apply sampling strategy to balance dataset"""
+    if strategy == "none":
+        return df
+    
+    label_counts = df['Label'].value_counts()
+    print(f"\nOriginal distribution: {dict(label_counts)}")
+    
+    if strategy == "balanced":
+        # Balanced: make both classes equal to average
+        target_count = int(label_counts.mean())
+        print(f"Target count per class: {target_count}")
+        
+    elif strategy == "oversample":
+        # Oversample minority to match majority
+        target_count = label_counts.max()
+        print(f"Target count per class (oversample): {target_count}")
+        
+    elif strategy == "undersample":
+        # Undersample majority to match minority
+        target_count = label_counts.min()
+        print(f"Target count per class (undersample): {target_count}")
+    
+    balanced_dfs = []
+    for label in [0, 1]:
+        label_df = df[df['Label'] == label]
+        current_count = len(label_df)
+        
+        if current_count < target_count:
+            # Oversample
+            sampled = label_df.sample(n=target_count, replace=True, random_state=42)
+        elif current_count > target_count:
+            # Undersample
+            sampled = label_df.sample(n=target_count, replace=False, random_state=42)
+        else:
+            sampled = label_df
+        
+        balanced_dfs.append(sampled)
+    
+    balanced_df = pd.concat(balanced_dfs, ignore_index=True)
+    balanced_df = balanced_df.sample(frac=1, random_state=42).reset_index(drop=True)  # Shuffle
+    
+    print(f"After sampling: {dict(balanced_df['Label'].value_counts())}")
+    return balanced_df
+
+if config.use_sampling:
+    print(f"\n4. Applying sampling strategy: {config.sampling_strategy}...")
+    train_df = apply_sampling(train_df, config.sampling_strategy)
+    print(f"Final train samples: {len(train_df)}")
+else:
+    print("\n4. Skipping sampling (using original distribution)...")
+
+# Load tokenizer
+print("\n5. Loading CodeBERT tokenizer...")
+tokenizer = RobertaTokenizer.from_pretrained(config.model_base_path)
+print(f"Tokenizer loaded: {tokenizer.__class__.__name__}")
+
+# Dataset class
+class WebAttackDataset(Dataset):
+    def __init__(self, dataframe, tokenizer, max_length):
+        self.data = dataframe.reset_index(drop=True)
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        text = str(self.data.loc[idx, 'Sentence'])
+        label = int(self.data.loc[idx, 'Label'])
+        
+        encoding = self.tokenizer(
+            text,
+            add_special_tokens=True,
+            max_length=self.max_length,
+            padding='max_length',
+            truncation=True,
+            return_tensors='pt'
+        )
+        
+        return {
+            'input_ids': encoding['input_ids'].flatten(),
+            'attention_mask': encoding['attention_mask'].flatten(),
+            'label': torch.tensor(label, dtype=torch.long)
+        }
+
+# Create datasets
+print("\n6. Creating datasets...")
+train_dataset = WebAttackDataset(train_df, tokenizer, config.max_length)
+test_dataset = WebAttackDataset(test_df, tokenizer, config.max_length)
+
+train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)
+test_loader = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False)
+
+print(f"Train batches: {len(train_loader)}")
+print(f"Test batches: {len(test_loader)}")
+
+# Model class
+class CodeBERTClassifier(nn.Module):
+    def __init__(self, model_path, num_labels=2, dropout=0.1):
+        super(CodeBERTClassifier, self).__init__()
+        self.codebert = RobertaModel.from_pretrained(model_path)
+        self.dropout = nn.Dropout(dropout)
+        self.classifier = nn.Linear(self.codebert.config.hidden_size, num_labels)
+    
+    def forward(self, input_ids, attention_mask):
+        outputs = self.codebert(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs.pooler_output
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+        return logits
+
+# Load model
+print("\n7. Loading CodeBERT model...")
+model = CodeBERTClassifier(config.model_base_path)
+model.to(config.device)
+print(f"Model loaded and moved to {config.device}")
+
+# Count parameters
+total_params = sum(p.numel() for p in model.parameters())
+trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+print(f"Total parameters: {total_params:,}")
+print(f"Trainable parameters: {trainable_params:,}")
+
+# Optimizer and scheduler
+optimizer = AdamW(model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay)
+total_steps = len(train_loader) * config.epochs
+scheduler = get_linear_schedule_with_warmup(
+    optimizer,
+    num_warmup_steps=config.warmup_steps,
+    num_training_steps=total_steps
+)
+
+criterion = nn.CrossEntropyLoss()
+
+# Training function
+def train_epoch(model, dataloader, optimizer, scheduler, criterion, device, gradient_accumulation_steps=4):
+    model.train()
+    total_loss = 0
+    predictions = []
+    true_labels = []
+    
+    optimizer.zero_grad()
+    
+    progress_bar = tqdm(dataloader, desc="Training")
+    for idx, batch in enumerate(progress_bar):
+        input_ids = batch['input_ids'].to(device)
+        attention_mask = batch['attention_mask'].to(device)
+        labels = batch['label'].to(device)
+        
+        logits = model(input_ids, attention_mask)
+        loss = criterion(logits, labels)
+        loss = loss / gradient_accumulation_steps  # Normalize loss
+        
+        loss.backward()
+        
+        if (idx + 1) % gradient_accumulation_steps == 0:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            scheduler.step()
+            optimizer.zero_grad()
+        
+        total_loss += loss.item() * gradient_accumulation_steps
+        
+        preds = torch.argmax(logits, dim=1)
+        predictions.extend(preds.cpu().numpy())
+        true_labels.extend(labels.cpu().numpy())
+        
+        progress_bar.set_postfix({'loss': loss.item() * gradient_accumulation_steps})
+    
+    avg_loss = total_loss / len(dataloader)
+    accuracy = accuracy_score(true_labels, predictions)
+    
+    return avg_loss, accuracy
+
+# Evaluation function
+def evaluate(model, dataloader, criterion, device):
+    model.eval()
+    total_loss = 0
+    predictions = []
+    true_labels = []
+    
+    with torch.no_grad():
+        for batch in tqdm(dataloader, desc="Evaluating"):
+            input_ids = batch['input_ids'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
+            labels = batch['label'].to(device)
+            
+            logits = model(input_ids, attention_mask)
+            loss = criterion(logits, labels)
+            
+            total_loss += loss.item()
+            
+            preds = torch.argmax(logits, dim=1)
+            predictions.extend(preds.cpu().numpy())
+            true_labels.extend(labels.cpu().numpy())
+    
+    avg_loss = total_loss / len(dataloader)
+    accuracy = accuracy_score(true_labels, predictions)
+    precision, recall, f1, _ = precision_recall_fscore_support(
+        true_labels, predictions, average='binary'
+    )
+    
+    return avg_loss, accuracy, precision, recall, f1, predictions, true_labels
+
+# Training loop
+print("\n8. Starting training...")
+print("="*80)
+
+best_accuracy = 0
+best_f1 = 0
+patience_counter = 0
+training_history = []
+
+for epoch in range(config.epochs):
+    print(f"\nEpoch {epoch + 1}/{config.epochs}")
+    print("-" * 80)
+    
+    # Train
+    train_loss, train_acc = train_epoch(
+        model, train_loader, optimizer, scheduler, criterion, config.device, config.gradient_accumulation_steps
+    )
+    
+    # Evaluate
+    test_loss, test_acc, test_precision, test_recall, test_f1, predictions, true_labels = evaluate(
+        model, test_loader, criterion, config.device
+    )
+    
+    # Log results
+    print(f"\nTrain Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}")
+    print(f"Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}")
+    print(f"Precision: {test_precision:.4f}, Recall: {test_recall:.4f}, F1: {test_f1:.4f}")
+    
+    # Save history
+    history = {
+        'epoch': epoch + 1,
+        'train_loss': train_loss,
+        'train_acc': train_acc,
+        'test_loss': test_loss,
+        'test_acc': test_acc,
+        'precision': test_precision,
+        'recall': test_recall,
+        'f1': test_f1
+    }
+    training_history.append(history)
+    
+    # Save best model
+    if test_f1 > best_f1:
+        best_f1 = test_f1
+        best_accuracy = test_acc
+        patience_counter = 0
+        
+        # Save PyTorch model
+        model_save_path = os.path.join(config.output_dir, 'best_model.pt')
+        torch.save({
+            'epoch': epoch + 1,
+            'model_state_dict': model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'test_acc': test_acc,
+            'test_f1': test_f1,
+            'config': vars(config)
+        }, model_save_path)
+        print(f"\n✓ Best model saved! (F1: {test_f1:.4f})")
+    else:
+        patience_counter += 1
+        print(f"\nNo improvement. Patience: {patience_counter}/{config.early_stopping_patience}")
+    
+    # Early stopping
+    if patience_counter >= config.early_stopping_patience:
+        print(f"\nEarly stopping triggered after {epoch + 1} epochs")
+        break
+
+print("\n" + "="*80)
+print("Training completed!")
+print("="*80)
+
+# Final evaluation
+print("\n9. Final evaluation on test set...")
+test_loss, test_acc, test_precision, test_recall, test_f1, predictions, true_labels = evaluate(
+    model, test_loader, criterion, config.device
+)
+
+print(f"\nFinal Test Results:")
+print(f"Accuracy: {test_acc:.4f}")
+print(f"Precision: {test_precision:.4f}")
+print(f"Recall: {test_recall:.4f}")
+print(f"F1 Score: {test_f1:.4f}")
+
+# Classification report
+print("\nClassification Report:")
+print(classification_report(true_labels, predictions, target_names=['Benign', 'Malicious']))
+
+# Confusion matrix
+cm = confusion_matrix(true_labels, predictions)
+print("\nConfusion Matrix:")
+print(cm)
+print(f"True Negatives: {cm[0][0]}")
+print(f"False Positives: {cm[0][1]}")
+print(f"False Negatives: {cm[1][0]}")
+print(f"True Positives: {cm[1][1]}")
+
+# Save results
+results = {
+    'final_metrics': {
+        'accuracy': float(test_acc),
+        'precision': float(test_precision),
+        'recall': float(test_recall),
+        'f1_score': float(test_f1)
+    },
+    'confusion_matrix': cm.tolist(),
+    'training_history': training_history,
+    'config': {
+        'epochs': config.epochs,
+        'batch_size': config.batch_size,
+        'learning_rate': config.learning_rate,
+        'max_length': config.max_length,
+        'sampling_strategy': config.sampling_strategy,
+        'train_samples': len(train_df),
+        'test_samples': len(test_df)
+    }
+}
+
+results_path = os.path.join(config.output_dir, 'training_results.json')
+with open(results_path, 'w') as f:
+    json.dump(results, f, indent=2)
+print(f"\nResults saved to: {results_path}")
+
+# Save tokenizer config
+tokenizer_config = {
+    'model_name': config.model_base_path,
+    'max_length': config.max_length
+}
+tokenizer_config_path = os.path.join(config.output_dir, 'tokenizer_config.json')
+with open(tokenizer_config_path, 'w') as f:
+    json.dump(tokenizer_config, f, indent=2)
+print(f"Tokenizer config saved to: {tokenizer_config_path}")
+
+print("\n" + "="*80)
+print("Training script completed successfully!")
+print(f"Best F1 Score: {best_f1:.4f}")
+print(f"Best Accuracy: {best_accuracy:.4f}")
+print(f"Model saved to: {config.output_dir}")
+print("="*80)

training_results.json

@@ -0,0 +1,59 @@
+{
+  "final_metrics": {
+    "accuracy": 0.993800976186482,
+    "precision": 0.9946941482577419,
+    "recall": 0.992129590365533,
+    "f1_score": 0.9934102141680395
+  },
+  "confusion_matrix": [
+    [
+      65914,
+      312
+    ],
+    [
+      464,
+      58491
+    ]
+  ],
+  "training_history": [
+    {
+      "epoch": 1,
+      "train_loss": 0.028937281491438017,
+      "train_acc": 0.9884227175957917,
+      "test_loss": 0.01923593317120965,
+      "test_acc": 0.9908692213674599,
+      "precision": 0.9826998864471311,
+      "recall": 0.9981850563989484,
+      "f1": 0.9903819453209806
+    },
+    {
+      "epoch": 2,
+      "train_loss": 0.020080875358571924,
+      "train_acc": 0.992353042207053,
+      "test_loss": 0.01691129035232433,
+      "test_acc": 0.9908053139054649,
+      "precision": 0.9840881682969316,
+      "recall": 0.9965906199643796,
+      "f1": 0.9902999351081672
+    },
+    {
+      "epoch": 3,
+      "train_loss": 0.01749216435263267,
+      "train_acc": 0.9929661568694804,
+      "test_loss": 0.02742091244277926,
+      "test_acc": 0.993800976186482,
+      "precision": 0.9946941482577419,
+      "recall": 0.992129590365533,
+      "f1": 0.9934102141680395
+    }
+  ],
+  "config": {
+    "epochs": 3,
+    "batch_size": 8,
+    "learning_rate": 2e-05,
+    "max_length": 256,
+    "sampling_strategy": "balanced",
+    "train_samples": 500722,
+    "test_samples": 125181
+  }
+}