---
language:
- en
- vi
tags:
- safety
- guardrail
- routing
- pytorch
- tabular-classification
metrics:
- f1
- accuracy
- precision
- recall
---

# SafeRoute Router Model (DynaGuard 1.7B / 8B)

This repository contains the weights for the **SafeRoute Router**, an optimized neural router designed to dynamically direct input prompts/responses between a lightweight safety classifier (Small Model) and a high-capacity safety classifier (Large Model). 

By routing "easy/safe" queries to the small model and reserving the large model only for "hard/unsafe" queries, the system drastically reduces inference latency and computational cost while preserving overall safety evaluation performance.

## Model Details

- **Architecture:** Multi-Layer Perceptron (MLP) with 3 hidden layers (`1024 -> 512 -> 256`), utilizing `BatchNorm1d`, `GELU` activations, and moderate `Dropout` (0.3).
- **Input Dimension:** `2048` (feature embeddings extracted from the small safety model).
- **Output Dimension:** `1` (binary classification logit indicating routing probability).
- **Loss Function:** `Focal Loss` ($\alpha=0.75, \gamma=2.0$) tailored to address severe class imbalance.
- **Optimizer & Scheduler:** `AdamW` with `CosineAnnealingWarmRestarts`.

## Evaluation Results

Evaluated on a balanced Test Benchmark at the optimal decision threshold (**0.6**):

| Metric | Score |
| :--- | :---: |
| **F1 Score** | **0.7525** |
| **Accuracy** | **0.7500** |
| **Precision** | **0.7451** |
| **Recall** | **0.7600** |
| **Overall AUPRC** | **0.7588** |

*Note: The high recall (0.76) combined with solid precision (0.74) ensures that potentially unsafe or ambiguous prompts are reliably intercepted and routed to the Large Model for thorough inspection.*

## How to Get Started with the Model

You can easily download and use this model in your PyTorch pipeline:

```python
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download

# 1. Define the Router Architecture
class RouterMLP(nn.Module):
    def __init__(self, input_dim=2048):
        super().__init__()
        self.cls = nn.Sequential(
            nn.Linear(input_dim, 1024),
            nn.BatchNorm1d(1024),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(1024, 512),
            nn.BatchNorm1d(512),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(256, 1),
        )

    def forward(self, x):
        return self.cls(x).squeeze(-1)

# 2. Download and Load the Checkpoint
repo_id = "YOUR_HF_USERNAME/safe-route-dynaguard" # <-- Replace with your repo name
model_path = hf_hub_download(repo_id=repo_id, filename="model.pt")

device = "cuda" if torch.cuda.is_available() else "cpu"
router = RouterMLP(input_dim=2048).to(device)

ckpt = torch.load(model_path, map_location=device)
router.load_state_dict(ckpt["state_dict"], strict=False)
router.eval()

# 3. Perform Routing Inference
with torch.no_grad():
    # Example feature tensor extracted from small model
    sample_features = torch.randn(4, 2048, device=device)
    
    logits = router(sample_features)
    routing_probs = torch.sigmoid(logits)
    
    # Use recommended threshold 0.6
    decisions = (routing_probs > 0.6).long()
    
    for i, decision in enumerate(decisions):
        if decision == 1:
            print(f"Sample {i}: Route to LARGE Model (Hard/Unsafe)")
        else:
            print(f"Sample {i}: Use SMALL Model (Easy/Safe)")
```

## Intended Use

- **Primary Use Case:** Guardrail optimization in LLM serving pipelines.
- **Out-of-Scope:** Standalone toxicity classification directly from raw text (this model requires intermediate hidden feature representations from a pre-trained small safety model).