e5-small-v2-hivetrace-finetune-CE-infoNCE

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 InfoNCE (contrastive) + cross-entropy on 32,538 triplets built from WildGuardMix train (plus HarmBench behaviors, about 1% of the data). Triplets include adversarial (jailbreak-style) prompts of both classes, jailbreak-wrapped safe negatives and nearest-neighbour "benign twins", so the model is pushed to judge content rather than form. Training: 2 epochs, batch 32, lr 2e-05, temperature 0.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-infoNCE"
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.7765   # 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-infoNCE")
embeddings = model.encode(["query: How do I reset my router?"])   # L2-normalized

Evaluation

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

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) and HarmBench behaviors. Evaluated on WildGuardTest and ToxicChat (lmsys/toxic-chat). Base model: intfloat/e5-small-v2 (MIT).