Namadgi/DeBERTav3-finetuned-threat-scorer

DeBERTa v3-base (183M params)
  → [CLS] token embedding
  → Dropout(0.1)
  → Linear(768, 1)
  → Sigmoid
  → threat score ∈ [0, 1]

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download


class ThreatRegressor(nn.Module):
    def __init__(self, model_name: str = "microsoft/deberta-v3-base", dropout: float = 0.1):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(model_name)
        hidden_size = self.encoder.config.hidden_size
        self.dropout = nn.Dropout(dropout)
        self.head = nn.Linear(hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls = out.last_hidden_state[:, 0, :]
        cls = self.dropout(cls)
        logit = self.head(cls).squeeze(-1)
        return logit

    @torch.no_grad()
    def predict_score(self, input_ids, attention_mask) -> torch.Tensor:
        return torch.sigmoid(self.forward(input_ids, attention_mask))


# Load model
model = ThreatRegressor("microsoft/deberta-v3-base", dropout=0.1)
weights_path = hf_hub_download("Namadgi/DeBERTav3-finetuned-threat-scorer", "model.safetensors")
state_dict = load_file(weights_path)
model.load_state_dict(state_dict)
model.eval()

# Tokenize and predict
tokenizer = AutoTokenizer.from_pretrained("Namadgi/DeBERTav3-finetuned-threat-scorer")
text = "I'm going to find you and hurt you"
enc = tokenizer(text, max_length=128, padding="max_length", truncation=True, return_tensors="pt")
score = model.predict_score(enc["input_ids"], enc["attention_mask"])
print(f"Threat score: {score.item():.4f}")  # e.g. 0.93 → 93% threat