Views
No views yet
1024 -> 512 -> 256), utilizing BatchNorm1d, GELU activations, and moderate Dropout (0.3).2048 (feature embeddings extracted from the small safety model).1 (binary classification logit indicating routing probability).Focal Loss ($\alpha=0.75, \gamma=2.0$) tailored to address severe class imbalance.AdamW with CosineAnnealingWarmRestarts.| Metric | Score |
|---|---|
| F1 Score | 0.7525 |
| Accuracy | 0.7500 |
| Precision | 0.7451 |
| Recall | 0.7600 |
| Overall AUPRC | 0.7588 |
1import torch
2import torch.nn as nn
3from huggingface_hub import hf_hub_download
4
5# 1. Define the Router Architecture
6class RouterMLP(nn.Module):
7 def __init__(self, input_dim=2048):
8 super().__init__()
9 self.cls = nn.Sequential(
10 nn.Linear(input_dim, 1024),
11 nn.BatchNorm1d(1024),
12 nn.GELU(),
13 nn.Dropout(0.3),
14 nn.Linear(1024, 512),
15 nn.BatchNorm1d(512),
16 nn.GELU(),
17 nn.Dropout(0.3),
18 nn.Linear(512, 256),
19 nn.BatchNorm1d(256),
20 nn.GELU(),
21 nn.Dropout(0.2),
22 nn.Linear(256, 1),
23 )
24
25 def forward(self, x):
26 return self.cls(x).squeeze(-1)
27
28# 2. Download and Load the Checkpoint
29repo_id = "StevenMup2004/DynaRoute" # <-- Replace with your repo name
30model_path = hf_hub_download(repo_id=repo_id, filename="model.pt")
31
32device = "cuda" if torch.cuda.is_available() else "cpu"
33router = RouterMLP(input_dim=2048).to(device)
34
35ckpt = torch.load(model_path, map_location=device)
36router.load_state_dict(ckpt["state_dict"], strict=False)
37router.eval()
38
39# 3. Perform Routing Inference
40with torch.no_grad():
41 # Example feature tensor extracted from small model
42 sample_features = torch.randn(4, 2048, device=device)
43
44 logits = router(sample_features)
45 routing_probs = torch.sigmoid(logits)
46
47 # Use recommended threshold 0.6
48 decisions = (routing_probs > 0.6).long()
49
50 for i, decision in enumerate(decisions):
51 if decision == 1:
52 print(f"Sample {i}: Route to LARGE Model (Hard/Unsafe)")
53 else:
54 print(f"Sample {i}: Use SMALL Model (Easy/Safe)")