initial mlx conversion of rampart pii ner

README.md

@@ -0,0 +1,87 @@
+---
+license: cc-by-4.0
+language:
+  - en
+library_name: mlx
+pipeline_tag: token-classification
+tags:
+  - mlx
+  - bert
+  - token-classification
+  - pii
+  - ner
+  - privacy
+base_model: nationaldesignstudio/rampart
+---
+
+# Rampart PII NER — MLX
+
+An MLX build of **Rampart**, a compact encoder-only BERT (MiniLM-L6, hidden 384,
+6 layers, ~18.5M params) with a 35-label BIO token-classification head for
+detecting personally identifiable information (PII). Intended for on-device,
+client-side PII redaction on Apple Silicon.
+
+This repository ships **float (fp16) MLX weights** in `model.safetensors` plus a
+small self-contained MLX implementation (`rampart_mlx.py`).
+
+## Provenance
+
+This is an **independent MLX conversion** of the original
+[`nationaldesignstudio/rampart`](https://huggingface.co/nationaldesignstudio/rampart).
+The original is distributed as a 4-bit quantized ONNX export; the float weights
+here were recovered directly from that export (4-bit `MatMulNBits` linears and
+INT8 embeddings dequantized to float) and then stored in MLX `safetensors`.
+
+The conversion was verified to reproduce the original ONNX model **exactly**:
+on the validation prompts, MLX vs. ONNX Runtime token-label agreement is
+**100%** with a maximum logit difference of ~1e-5 (floating-point rounding).
+
+No third-party MLX port was used in producing these weights.
+
+## Labels
+
+17 entity types in BIO format (35 classes incl. `O`): `GIVEN_NAME`, `SURNAME`,
+`EMAIL`, `PHONE`, `URL`, `TAX_ID`, `BANK_ACCOUNT`, `ROUTING_NUMBER`,
+`GOVERNMENT_ID`, `PASSPORT`, `DRIVERS_LICENSE`, `BUILDING_NUMBER`, `STREET_NAME`,
+`SECONDARY_ADDRESS`, `CITY`, `STATE`, `ZIP_CODE`.
+
+## Usage
+
+```bash
+pip install mlx tokenizers
+python demo.py "My name is John Smith, email john.smith@example.com"
+```
+
+```python
+import mlx.core as mx
+from tokenizers import Tokenizer
+from rampart_mlx import load
+
+model, cfg = load(".")
+tok = Tokenizer.from_file("tokenizer.json")
+enc = tok.encode("Call me at (555) 123-4567")
+logits = model(mx.array([enc.ids]), mx.array([enc.attention_mask]))
+label_ids = mx.argmax(logits[0], axis=-1).tolist()
+labels = [cfg.id2label[i] for i in label_ids]
+```
+
+See `demo.py` for BIO span decoding using the tokenizer's char offsets (needed to
+map predicted labels back onto the original text for redaction).
+
+## Files
+
+| File | Purpose |
+|------|---------|
+| `model.safetensors` | fp16 MLX weights (HuggingFace-style key names) |
+| `config.json` | model architecture + `id2label` |
+| `rampart_mlx.py` | self-contained MLX model + loader |
+| `demo.py` | tokenize → infer → decode spans |
+| `tokenizer.json`, `vocab.txt`, `tokenizer_config.json`, `special_tokens_map.json` | WordPiece tokenizer |
+
+## License & attribution
+
+Released under **CC-BY-4.0**, the same license as the upstream model. Attribution:
+
+- Original model: [`nationaldesignstudio/rampart`](https://huggingface.co/nationaldesignstudio/rampart) (CC-BY-4.0)
+- Base encoder: [`nreimers/MiniLM-L6-H384-uncased`](https://huggingface.co/nreimers/MiniLM-L6-H384-uncased)
+- Training data: [`ai4privacy/pii-masking-openpii-1.5m`](https://huggingface.co/datasets/ai4privacy/pii-masking-openpii-1.5m) (CC-BY-4.0)

config.json

@@ -0,0 +1,99 @@
+{
+  "architectures": [
+    "BertForTokenClassification"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "dtype": "float32",
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 384,
+  "id2label": {
+    "0": "O",
+    "1": "B-GIVEN_NAME",
+    "2": "I-GIVEN_NAME",
+    "3": "B-SURNAME",
+    "4": "I-SURNAME",
+    "5": "B-EMAIL",
+    "6": "I-EMAIL",
+    "7": "B-PHONE",
+    "8": "I-PHONE",
+    "9": "B-URL",
+    "10": "I-URL",
+    "11": "B-TAX_ID",
+    "12": "I-TAX_ID",
+    "13": "B-BANK_ACCOUNT",
+    "14": "I-BANK_ACCOUNT",
+    "15": "B-ROUTING_NUMBER",
+    "16": "I-ROUTING_NUMBER",
+    "17": "B-GOVERNMENT_ID",
+    "18": "I-GOVERNMENT_ID",
+    "19": "B-PASSPORT",
+    "20": "I-PASSPORT",
+    "21": "B-DRIVERS_LICENSE",
+    "22": "I-DRIVERS_LICENSE",
+    "23": "B-BUILDING_NUMBER",
+    "24": "I-BUILDING_NUMBER",
+    "25": "B-STREET_NAME",
+    "26": "I-STREET_NAME",
+    "27": "B-SECONDARY_ADDRESS",
+    "28": "I-SECONDARY_ADDRESS",
+    "29": "B-CITY",
+    "30": "I-CITY",
+    "31": "B-STATE",
+    "32": "I-STATE",
+    "33": "B-ZIP_CODE",
+    "34": "I-ZIP_CODE"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "label2id": {
+    "B-BANK_ACCOUNT": 13,
+    "B-BUILDING_NUMBER": 23,
+    "B-CITY": 29,
+    "B-DRIVERS_LICENSE": 21,
+    "B-EMAIL": 5,
+    "B-GIVEN_NAME": 1,
+    "B-GOVERNMENT_ID": 17,
+    "B-PASSPORT": 19,
+    "B-PHONE": 7,
+    "B-ROUTING_NUMBER": 15,
+    "B-SECONDARY_ADDRESS": 27,
+    "B-STATE": 31,
+    "B-STREET_NAME": 25,
+    "B-SURNAME": 3,
+    "B-TAX_ID": 11,
+    "B-URL": 9,
+    "B-ZIP_CODE": 33,
+    "I-BANK_ACCOUNT": 14,
+    "I-BUILDING_NUMBER": 24,
+    "I-CITY": 30,
+    "I-DRIVERS_LICENSE": 22,
+    "I-EMAIL": 6,
+    "I-GIVEN_NAME": 2,
+    "I-GOVERNMENT_ID": 18,
+    "I-PASSPORT": 20,
+    "I-PHONE": 8,
+    "I-ROUTING_NUMBER": 16,
+    "I-SECONDARY_ADDRESS": 28,
+    "I-STATE": 32,
+    "I-STREET_NAME": 26,
+    "I-SURNAME": 4,
+    "I-TAX_ID": 12,
+    "I-URL": 10,
+    "I-ZIP_CODE": 34,
+    "O": 0
+  },
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 6,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "transformers_version": "4.57.6",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 19730
+}

demo.py

@@ -0,0 +1,55 @@
+"""
+Minimal Rampart-MLX demo: tokenize, run the model, decode BIO spans.
+
+    python demo.py "My name is John Smith, email john@example.com"
+"""
+
+import sys
+
+import mlx.core as mx
+from tokenizers import Tokenizer
+
+from rampart_mlx import load
+
+
+def decode_spans(tokens, offsets, label_ids, id2label):
+    """Merge B-/I- runs into (entity_type, text, start, end) spans."""
+    spans, cur = [], None
+    for (start, end), lid in zip(offsets, label_ids):
+        if start == end:  # special token
+            continue
+        label = id2label[lid]
+        if label == "O":
+            cur = None
+            continue
+        tag, etype = label[0], label[2:]
+        if tag == "B" or cur is None or cur["type"] != etype:
+            cur = {"type": etype, "start": start, "end": end}
+            spans.append(cur)
+        else:
+            cur["end"] = end
+    return spans
+
+
+def main():
+    text = sys.argv[1] if len(sys.argv) > 1 else \
+        "My name is John Smith and my email is john.smith@example.com"
+    model, cfg = load(".")
+    tok = Tokenizer.from_file("tokenizer.json")
+    enc = tok.encode(text)
+
+    input_ids = mx.array([enc.ids])
+    attention_mask = mx.array([enc.attention_mask])
+    logits = model(input_ids, attention_mask)
+    label_ids = mx.argmax(logits[0], axis=-1).tolist()
+
+    spans = decode_spans(enc.tokens, enc.offsets, label_ids, cfg.id2label)
+    print(f"\nText: {text}\n")
+    if not spans:
+        print("No PII detected.")
+    for s in spans:
+        print(f"  {s['type']:16s} {text[s['start']:s['end']]!r}  [{s['start']}:{s['end']}]")
+
+
+if __name__ == "__main__":
+    main()

model.safetensors

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

rampart_mlx.py

@@ -0,0 +1,132 @@
+"""
+Rampart PII NER — MLX implementation.
+
+A standard encoder-only BERT (MiniLM-L6, hidden 384) with a 35-label BIO
+token-classification head for PII detection. Weights are float (fp16) and load
+directly from `model.safetensors`.
+
+Independent MLX conversion: float weights were recovered 1:1 from the original
+`nationaldesignstudio/rampart` 4-bit ONNX export (verified to match the ONNX
+runtime output exactly). See README for attribution.
+"""
+
+import json
+import math
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class BertConfig:
+    def __init__(self, path):
+        d = json.loads(Path(path).read_text())
+        self.hidden_size = d["hidden_size"]
+        self.num_hidden_layers = d["num_hidden_layers"]
+        self.num_attention_heads = d["num_attention_heads"]
+        self.intermediate_size = d["intermediate_size"]
+        self.vocab_size = d["vocab_size"]
+        self.max_position_embeddings = d["max_position_embeddings"]
+        self.type_vocab_size = d["type_vocab_size"]
+        self.layer_norm_eps = d["layer_norm_eps"]
+        self.id2label = {int(k): v for k, v in d["id2label"].items()}
+        self.num_labels = len(self.id2label)
+
+
+class BertLayer(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        H, I = c.hidden_size, c.intermediate_size
+        self.query = nn.Linear(H, H)
+        self.key = nn.Linear(H, H)
+        self.value = nn.Linear(H, H)
+        self.attn_out = nn.Linear(H, H)
+        self.attn_ln = nn.LayerNorm(H, eps=c.layer_norm_eps)
+        self.intermediate = nn.Linear(H, I)
+        self.output = nn.Linear(I, H)
+        self.output_ln = nn.LayerNorm(H, eps=c.layer_norm_eps)
+        self.n_heads = c.num_attention_heads
+        self.d_head = H // c.num_attention_heads
+
+    def __call__(self, x, add_mask):
+        B, S, _ = x.shape
+
+        def split(t):
+            return t.reshape(B, S, self.n_heads, self.d_head).transpose(0, 2, 1, 3)
+
+        q, k, v = split(self.query(x)), split(self.key(x)), split(self.value(x))
+        scores = (q @ k.transpose(0, 1, 3, 2)) / math.sqrt(self.d_head) + add_mask
+        ctx = (mx.softmax(scores, axis=-1) @ v).transpose(0, 2, 1, 3).reshape(B, S, -1)
+        x = self.attn_ln(self.attn_out(ctx) + x)
+        h = nn.gelu(self.intermediate(x))
+        return self.output_ln(self.output(h) + x)
+
+
+class RampartForTokenClassification(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.config = c
+        self.word_embeddings = nn.Embedding(c.vocab_size, c.hidden_size)
+        self.position_embeddings = nn.Embedding(c.max_position_embeddings, c.hidden_size)
+        self.token_type_embeddings = nn.Embedding(c.type_vocab_size, c.hidden_size)
+        self.embeddings_ln = nn.LayerNorm(c.hidden_size, eps=c.layer_norm_eps)
+        self.layers = [BertLayer(c) for _ in range(c.num_hidden_layers)]
+        self.classifier = nn.Linear(c.hidden_size, c.num_labels)
+
+    def __call__(self, input_ids, attention_mask, token_type_ids=None):
+        B, S = input_ids.shape
+        if token_type_ids is None:
+            token_type_ids = mx.zeros((B, S), dtype=input_ids.dtype)
+        pos_ids = mx.arange(S)[None, :]
+        x = self.embeddings_ln(
+            self.word_embeddings(input_ids)
+            + self.position_embeddings(pos_ids)
+            + self.token_type_embeddings(token_type_ids)
+        )
+        add_mask = (1.0 - attention_mask.astype(mx.float32))[:, None, None, :] * -1e9
+        for layer in self.layers:
+            x = layer(x, add_mask)
+        return self.classifier(x)
+
+
+# Maps the HuggingFace-style safetensors keys onto this module tree.
+def _remap(weights, n_layers):
+    out = {}
+    direct = {
+        "bert.embeddings.word_embeddings.weight": "word_embeddings.weight",
+        "bert.embeddings.position_embeddings.weight": "position_embeddings.weight",
+        "bert.embeddings.token_type_embeddings.weight": "token_type_embeddings.weight",
+        "bert.embeddings.LayerNorm.weight": "embeddings_ln.weight",
+        "bert.embeddings.LayerNorm.bias": "embeddings_ln.bias",
+        "classifier.weight": "classifier.weight",
+        "classifier.bias": "classifier.bias",
+    }
+    for src, dst in direct.items():
+        out[dst] = weights[src]
+    for i in range(n_layers):
+        p = f"bert.encoder.layer.{i}"
+        m = {
+            f"{p}.attention.self.query": f"layers.{i}.query",
+            f"{p}.attention.self.key": f"layers.{i}.key",
+            f"{p}.attention.self.value": f"layers.{i}.value",
+            f"{p}.attention.output.dense": f"layers.{i}.attn_out",
+            f"{p}.attention.output.LayerNorm": f"layers.{i}.attn_ln",
+            f"{p}.intermediate.dense": f"layers.{i}.intermediate",
+            f"{p}.output.dense": f"layers.{i}.output",
+            f"{p}.output.LayerNorm": f"layers.{i}.output_ln",
+        }
+        for src, dst in m.items():
+            out[f"{dst}.weight"] = weights[f"{src}.weight"]
+            out[f"{dst}.bias"] = weights[f"{src}.bias"]
+    return out
+
+
+def load(model_dir):
+    model_dir = Path(model_dir)
+    c = BertConfig(model_dir / "config.json")
+    model = RampartForTokenClassification(c)
+    weights = mx.load(str(model_dir / "model.safetensors"))
+    flat = _remap(weights, c.num_hidden_layers)
+    model.load_weights(list(flat.items()))
+    mx.eval(model.parameters())
+    return model, c

special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,58 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": true,
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "never_split": null,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}