README.md

---
license: mit
language:
- en
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