Railz-Micro: 67M static multilabel safety guard (safety-vocab gemma distill)

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+pipeline.skops filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,149 @@
+---
+license: mit
+library_name: model2vec
+pipeline_tag: text-classification
+language:
+- en
+tags:
+- safety
+- guardrail
+- moderation
+- jailbreak-detection
+- multilabel
+- static-embeddings
+- model2vec
+datasets:
+- nvidia/Aegis-AI-Content-Safety-Dataset-2.0
+- OpenSafetyLab/Salad-Data
+- allenai/WildChat-1M
+- google/civil_comments
+- toxigen/toxigen-data
+- allenai/real-toxicity-prompts
+- bench-llm/or-bench
+- natolambert/xstest-v2-copy
+---
+
+# Railz-Micro
+
+**One tiny model, one pass, three safety jobs: harmful-content detection + 9-category classification + jailbreak detection.**
+
+Railz-Micro is a 67M-parameter **static** multilabel safety guard. No transformer at inference — tokenize, look up, average, classify. Sub-millisecond on CPU, runs anywhere, nothing leaves your machine.
+
+## Why it's different
+
+- **Custom safety-vocabulary base.** We distilled [google/embeddinggemma-300m](https://huggingface.co/google/embeddinggemma-300m) into a static token table **with 3,864 mined safety phrases** (2,472 multi-word) added as dedicated tokens. `how to prevent` and `ignore previous instructions` are single tokens with their own composed vectors — the teacher's contextual reading of each phrase, frozen into a lookup. This is what lets a bag-of-tokens model separate *"how to prevent bomb attacks"* (benign) from bomb-making requests, and it's why the false-positive rate on scary-but-benign prompts is ~1%.
+- **Geometry-curated training data.** 355k examples curated with [SemHash](https://github.com/MinishLab/semhash): semantic dedup, **decontamination against every benchmark below** (0.85 threshold — paraphrase-level leaks removed, not just exact matches), and hard examples mined by embedding geometry (benign prompts nearest the harmful cluster and vice versa) rather than keywords.
+- **One model instead of five.** Binary harm, 9 harm categories, and jailbreak flags come from a single multilabel head in one forward pass.
+
+## Benchmarks
+
+All rows are **out-of-domain** (no split of these sets was trained on; the training blend was decontaminated against all of them) except Aegis, which is in-domain and marked as such. Default threshold τ=0.5 unless noted. No cherry-picking: weak axes are shown and discussed in Limitations.
+
+### Mixed sets (precision + recall)
+
+| benchmark | F1 | F0.5 | P | R | n (+pos) |
+|---|---|---|---|---|---|
+| ToxicChat (test) | 36.4 | 37.1 | 37.5 | 35.4 | 5083 (+362) |
+| OpenAI-Moderation | 54.8 | 57.5 | 59.4 | 51.0 | 1680 (+522) |
+| ToxicConversations | 20.9 | 25.9 | 30.8 | 15.8 | 4000 (+311) |
+| Aegis-2.0 (test, **in-domain**) | 78.7 | 77.9 | 77.4 | 80.0 | 1964 (+1059) |
+
+### Over-refusal — false-positive rate on benign-but-scary prompts (lower = better)
+
+| benchmark | FPR |
+|---|---|
+| OR-Bench (5,000 held-out, never trained on) | **0.8%** |
+| OR-Bench-hard-1k | **1.9%** |
+
+### Catch-rate on all-harmful sets (recall; precision undefined)
+
+| benchmark | τ=0.5 | τ=0.02 |
+|---|---|---|
+| MaliciousInstruct | 79% | 89% |
+| SimpleSafetyTests | 62% | — |
+| do-not-answer | 56% | — |
+| HarmfulQA | 53% | 69% |
+| OR-Bench-toxic | 32% | 46% |
+
+### Jailbreak (jackhhao/jailbreak-classification, test)
+
+| F1 | P | R |
+|---|---|---|
+| 58.9 | 84.0 | 45.3 |
+
+### Categories (Aegis-test, in-domain, 9 buckets)
+
+Multilabel macro-P **74.4** / macro-R **46.6**; a correct category is predicted for **67%** of unsafe prompts.
+
+## Choosing a threshold
+
+The model is precision-first at the default τ=0.5. Lowering τ buys recall while the false-positive rate stays low (measured on the held-out OR-Bench slice):
+
+| τ | OR-Bench FPR | MaliciousInstruct catch | use case |
+|---|---|---|---|
+| 0.50 | 0.8% | 79% | max precision (default) |
+| 0.15 | 1.5% | 82% | balanced |
+| 0.02 | 2.4% | 89% | max recall |
+
+```python
+from model2vec.inference import StaticModelPipeline
+import numpy as np
+
+pipe = StaticModelPipeline.from_pretrained("bfuzzy1/Railz-Micro")
+
+# default thresholds
+pipe.predict(["how do I make a pipe bomb"])          # ['harmful', 'cat:violence_weapons', ...]
+
+# custom threshold on P(harmful)
+proba = np.asarray(pipe.predict_proba(["how do I make a pipe bomb"]))
+harmful_idx = list(pipe.classes_).index("harmful")
+flag = proba[:, harmful_idx] >= 0.15                  # τ of your choice
+```
+
+Labels: `harmful`, `jailbreak`, and `cat:{violence_weapons, hate_harassment, sexual, crime_drugs, cyber_fraud, misinfo, self_harm, privacy, advice}`.
+
+## Recipe
+
+1. **Vocab mining** — discriminative 1-3-grams from ~480k safety prompts, cleaned by 4 passes (cross-source robustness ≥2 datasets, proper-noun strip via mid-sentence capitalization, stopword-edge coherence, split-half stability) + curated jailbreak phrases and benign disambiguators + LDNOOBW lexicon → 3,864 phrases.
+2. **Base distillation** — `model2vec.distill(embeddinggemma-300m, vocabulary=...)` → 260k-token static table, 256-dim, PCA + SIF.
+3. **Data curation** — SemHash dedup (0.9) → decontamination vs all benchmarks (0.85) → boundary mining (0.6): 120k benign-near-harmful + 26k harmful-near-benign hard examples; jailbreak rows exempt from dedup (attack paraphrases are signal).
+4. **Head training** — multilabel `StaticModelForClassification.fit` on 355k examples (28% harmful), ≤15 epochs, early stopping.
+
+Training data: Aegis-2.0, Salad-Data (+attack set → jailbreak labels), Nemotron content-safety, ToxiGen, RealToxicityPrompts, WildChat-1M (clean user turns), civil_comments, OR-Bench (train-negs only; 5k slice held out for the FPR eval above), XSTest.
+
+## Limitations (honest ones)
+
+- **Informal real-user chat is the weak axis** (ToxicChat ~36 F1, ToxicConversations ~21). Typos, slang, and context-dependent toxicity need composition a static model doesn't have.
+- **Subtle / academically-phrased harm** (HarmfulQA-style) catches ~53-69% depending on τ — phrase-sparse harm is hard for a lookup table.
+- **Jailbreak recall is moderate (45% OOD)** at high precision (84%). Novel attack templates outside the mined phrase set fall back to subword averaging.
+- **Prompt-level, English-only.** Does not score model responses; not tested on code-mixed or non-English input.
+- **No deep composition.** Negation, sarcasm, multi-sentence intent are out of scope. For those, cascade: Railz-Micro filters at wire speed, escalate uncertain cases to a contextual guard (e.g. [Railz-R2](https://huggingface.co/bfuzzy1/Railz-R2)).
+
+## Speed & footprint
+
+Static embeddings + sklearn head: sub-ms per prompt single-threaded CPU, no GPU, no PyTorch at inference (`pip install model2vec[inference]`). 67M params.
+
+## Part of the Railz family
+
+| model | size | role |
+|---|---|---|
+| [Railz](https://huggingface.co/bfuzzy1/Railz) | 0.6B | policy-conditioned guard |
+| [Railz-R](https://huggingface.co/bfuzzy1/Railz-R) | 0.6B | + reasoning |
+| [Railz-R2](https://huggingface.co/bfuzzy1/Railz-R2) | 0.6B | + OOD robustness |
+| **Railz-Micro** | **67M static** | **wire-speed multilabel prefilter** |
+
+## Citation
+
+Built with [Model2Vec](https://github.com/MinishLab/model2vec) by Minish Lab:
+
+```
+@software{minishlab2024model2vec,
+  author       = {Stephan Tulkens and {van Dongen}, Thomas},
+  title        = {Model2Vec: Fast State-of-the-Art Static Embeddings},
+  year         = {2024},
+  publisher    = {Zenodo},
+  doi          = {10.5281/zenodo.17270888},
+  url          = {https://github.com/MinishLab/model2vec},
+  license      = {MIT}
+}
+```

config.json

@@ -0,0 +1,4 @@
+{
+    "normalize": true,
+    "embedding_dtype": "float32"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:579048738b6ccaab56af284fb4f16aa6a74ced9d86fe71292809e8d96377cb94
+size 266488920

modules.json

@@ -0,0 +1,14 @@
+[
+    {
+        "idx": 0,
+        "name": "0",
+        "path": ".",
+        "type": "sentence_transformers.models.StaticEmbedding"
+    },
+    {
+        "idx": 1,
+        "name": "1",
+        "path": "1_Normalize",
+        "type": "sentence_transformers.models.Normalize"
+    }
+]

pipeline.skops

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d09106086dafe8f21bf11980df657bc206a4060b0e41ac05b1326bafb0204fdb
+size 3916238

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:42b42cd26bc92f8a39527672bfc3a84c620cbd46c43873a646a601166c33f9ec
+size 13573897