e5-small-v2-hivetrace-finetune-CE

A compact encoder intended as a fast pre-filter in front of an LLM: it flags likely harmful user prompts so that only suspicious ones go to a heavier safety check.

intfloat/e5-small-v2 fully fine-tuned with cross-entropy (linear head) on 43,025 labeled prompts from WildGuardMix train (deduplicated). Only real WildGuardMix prompts are used: no synthetic or template-generated data. Training: 2 epochs, batch 32, lr 2e-05, single seed.

Usage

Every input must be prefixed with "query: ", as in e5.

Classifier (encoder + head). The head is a single Linear(384, 2) trained on the raw, un-normalized mean-pooled embedding: do not L2-normalize before the head.

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

repo = "pineee/e5-small-v2-hivetrace-finetune-CE"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModel.from_pretrained(repo).eval()
head = torch.nn.Linear(384, 2)
head.load_state_dict(load_file(hf_hub_download(repo, "head.safetensors")))

def embed(texts):
    batch = tokenizer(["query: " + t for t in texts], padding=True, truncation=True,
                      max_length=512, return_tensors="pt")
    with torch.no_grad():
        hidden = model(**batch).last_hidden_state
    mask = batch["attention_mask"].unsqueeze(-1).float()
    return (hidden * mask).sum(1) / mask.sum(1)   # raw mean pooling, no normalization

p_harm = torch.softmax(head(embed(["How do I reset my router?"])), dim=-1)[:, 1]
flagged = p_harm >= 0.9711   # 1% false positives on WildGuardMix validation

Embeddings only (similarity search, your own classifier on top):

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("pineee/e5-small-v2-hivetrace-finetune-CE")
embeddings = model.encode(["query: How do I reset my router?"])   # L2-normalized

Evaluation

Head with threshold 0.9711, set at 1% false positive rate on WildGuardMix validation and not re-tuned on the test sets:

set	n	AUROC	TPR@FPR=1%	FPR at threshold	FNR at threshold
WildGuardTest (in-domain)	1,699	0.925	0.398	2.2%	39.0%
ToxicChat test (out-of-domain)	2,853	0.925	0.122	5.1%	32.9%

At a false positive rate of a few percent the model still misses a large share of harmful prompts (FNR column).

Limitations

• English only, single training seed.
• WildGuardMix train labels are produced by an LLM and contain noise.
• Calibration drifts across data sources: re-calibrate the threshold on your own traffic.
• Use as a first-stage filter, not as the only safeguard.

Data and attribution

Trained on WildGuardMix (allenai/wildguardmix, ODC-By). Evaluated on WildGuardTest and ToxicChat (lmsys/toxic-chat). Base model: intfloat/e5-small-v2 (MIT).