Multi-task: PII NER (F1=0.49) + 10-class doc clf (acc=0.25)

.gitattributes

@@ -33,3 +33,4 @@ 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
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,126 @@
+---
+license: apache-2.0
+base_model: openai/privacy-filter
+tags:
+- token-classification
+- text-classification
+- multi-task
+- pii-detection
+- document-classification
+- privacy
+datasets:
+- ai4privacy/pii-masking-400k
+- community-datasets/yahoo_answers_topics
+metrics:
+- f1
+- accuracy
+model-index:
+- name: privacy-filter-multitask
+  results:
+  - task:
+      type: token-classification
+      name: PII Detection
+    dataset:
+      name: ai4privacy/pii-masking-400k
+      type: ai4privacy/pii-masking-400k
+    metrics:
+    - type: f1
+      value: 0.4925
+    - type: precision
+      value: 0.6968
+    - type: recall
+      value: 0.3809
+  - task:
+      type: text-classification
+      name: Document Classification
+    dataset:
+      name: yahoo_answers_topics
+      type: community-datasets/yahoo_answers_topics
+    metrics:
+    - type: accuracy
+      value: 0.2482
+---
+
+# Privacy Filter Multi-Task 🔒📄
+
+A **single model** for simultaneous **PII Detection (NER)** and **Document Classification (10 categories)**.
+
+Adapted from [openai/privacy-filter](https://huggingface.co/openai/privacy-filter) (1.4B Sparse MoE, ~50M active params/token).
+
+## Architecture
+
+```
+Input → BPE Tokenizer (200K vocab)
+  ↓
+8-layer Sparse MoE Transformer (128 experts, top-4)
+  ↓                      ↓
+NER Head (640→33)    Doc Head (mean-pool → 640→10)
+  ↓                      ↓
+BIOES PII tags        10 categories
+```
+
+## Results
+
+| Task | Metric | Value |
+|------|--------|-------|
+| PII NER | F1 (strict, span) | **0.493** |
+| PII NER | Precision | 0.697 |
+| PII NER | Recall | 0.381 |
+| PII NER | Token Accuracy | 0.944 |
+| Doc Clf | Val Accuracy | 0.255 |
+| Doc Clf | Test Accuracy | **0.248** |
+
+### Inference Speed
+| Device | Latency |
+|--------|---------|
+| GPU A10G (bf16) | 178 ms |
+| CPU (fp32) | 202 ms |
+
+## PII Entity Types
+`private_person` • `private_email` • `private_phone` • `private_address` • `private_date` • `private_url` • `account_number` • `secret`
+
+## Document Categories
+Society & Culture • Science & Math • Health • Education • Computers & Internet • Sports • Business & Finance • Entertainment • Family • Politics
+
+## Usage
+
+```python
+import torch, torch.nn as nn
+from transformers import AutoModelForTokenClassification, AutoTokenizer
+from huggingface_hub import hf_hub_download
+
+tokenizer = AutoTokenizer.from_pretrained("binga/privacy-filter-multitask")
+model = AutoModelForTokenClassification.from_pretrained(
+    "binga/privacy-filter-multitask", dtype=torch.bfloat16, device_map="auto"
+)
+doc_head = nn.Linear(640, 10)
+doc_head.load_state_dict(torch.load(
+    hf_hub_download("binga/privacy-filter-multitask", "doc_head.pt"),
+    weights_only=True, map_location=model.device
+))
+doc_head = doc_head.to(dtype=torch.bfloat16, device=model.device)
+
+text = "John Smith (SSN: 123-45-6789) emailed john@corp.com"
+inputs = tokenizer(text, return_tensors="pt").to(model.device)
+with torch.no_grad():
+    out = model(**inputs, output_hidden_states=True)
+
+# PII
+for t, p in zip(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]),
+                out.logits.argmax(-1)[0]):
+    label = model.config.id2label[p.item()]
+    if label != "O": print(f"  {t} → {label}")
+
+# Doc class
+cats = ["Society", "Science", "Health", "Education", "Computers",
+        "Sports", "Business", "Entertainment", "Family", "Politics"]
+h = out.hidden_states[-1]
+m = inputs["attention_mask"].unsqueeze(-1).to(h.dtype)
+pooled = (h * m).sum(1) / m.sum(1).clamp(min=1)
+print(f"Category: {cats[doc_head(pooled).argmax().item()]}")
+```
+
+## Training
+- Partial fine-tune: last 4/8 MoE layers + heads (636M/1.4B trainable)
+- NER: ai4privacy/pii-masking-400k (20K en), Doc: yahoo_answers_topics (20K)
+- Loss: NER×1.0 + Doc×0.5, AdamW LR=2e-5, cosine, 2 epochs, BS=16

config.json

@@ -0,0 +1,120 @@
+{
+  "architectures": [
+    "OpenAIPrivacyFilterForTokenClassification"
+  ],
+  "attention_bias": true,
+  "attention_dropout": 0.0,
+  "bos_token_id": null,
+  "classifier_dropout": 0.0,
+  "default_n_ctx": 128000,
+  "dtype": "bfloat16",
+  "eos_token_id": 199999,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 640,
+  "id2label": {
+    "0": "O",
+    "1": "B-account_number",
+    "2": "I-account_number",
+    "3": "E-account_number",
+    "4": "S-account_number",
+    "5": "B-private_address",
+    "6": "I-private_address",
+    "7": "E-private_address",
+    "8": "S-private_address",
+    "9": "B-private_date",
+    "10": "I-private_date",
+    "11": "E-private_date",
+    "12": "S-private_date",
+    "13": "B-private_email",
+    "14": "I-private_email",
+    "15": "E-private_email",
+    "16": "S-private_email",
+    "17": "B-private_person",
+    "18": "I-private_person",
+    "19": "E-private_person",
+    "20": "S-private_person",
+    "21": "B-private_phone",
+    "22": "I-private_phone",
+    "23": "E-private_phone",
+    "24": "S-private_phone",
+    "25": "B-private_url",
+    "26": "I-private_url",
+    "27": "E-private_url",
+    "28": "S-private_url",
+    "29": "B-secret",
+    "30": "I-secret",
+    "31": "E-secret",
+    "32": "S-secret"
+  },
+  "initial_context_length": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 640,
+  "label2id": {
+    "B-account_number": 1,
+    "B-private_address": 5,
+    "B-private_date": 9,
+    "B-private_email": 13,
+    "B-private_person": 17,
+    "B-private_phone": 21,
+    "B-private_url": 25,
+    "B-secret": 29,
+    "E-account_number": 3,
+    "E-private_address": 7,
+    "E-private_date": 11,
+    "E-private_email": 15,
+    "E-private_person": 19,
+    "E-private_phone": 23,
+    "E-private_url": 27,
+    "E-secret": 31,
+    "I-account_number": 2,
+    "I-private_address": 6,
+    "I-private_date": 10,
+    "I-private_email": 14,
+    "I-private_person": 18,
+    "I-private_phone": 22,
+    "I-private_url": 26,
+    "I-secret": 30,
+    "O": 0,
+    "S-account_number": 4,
+    "S-private_address": 8,
+    "S-private_date": 12,
+    "S-private_email": 16,
+    "S-private_person": 20,
+    "S-private_phone": 24,
+    "S-private_url": 28,
+    "S-secret": 32
+  },
+  "max_position_embeddings": 131072,
+  "model_type": "openai_privacy_filter",
+  "num_attention_heads": 14,
+  "num_experts_per_tok": 4,
+  "num_hidden_layers": 8,
+  "num_key_value_heads": 2,
+  "num_local_experts": 128,
+  "output_router_logits": false,
+  "pad_token_id": 199999,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 32.0,
+    "original_max_position_embeddings": 4096,
+    "rope_theta": 150000.0,
+    "rope_type": "yarn",
+    "truncate": false
+  },
+  "router_aux_loss_coef": 0.001,
+  "sliding_window": 128,
+  "tie_word_embeddings": false,
+  "transformers.js_config": {
+    "use_external_data_format": {
+      "model": 1,
+      "model.onnx": 3,
+      "model_fp16.onnx": 2
+    }
+  },
+  "transformers_version": "5.6.2",
+  "use_cache": false,
+  "vocab_size": 200064
+}

doc_head.pt

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

model.safetensors

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

multitask_config.json

@@ -0,0 +1,17 @@
+{
+  "num_doc_classes": 10,
+  "doc_label_names": [
+    "Society & Culture",
+    "Science & Mathematics",
+    "Health",
+    "Education & Reference",
+    "Computers & Internet",
+    "Sports",
+    "Business & Finance",
+    "Entertainment & Music",
+    "Family & Relationships",
+    "Politics & Government"
+  ],
+  "ner_weight": 1.0,
+  "doc_weight": 0.5
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "backend": "tokenizers",
+  "eos_token": "<|endoftext|>",
+  "is_local": true,
+  "local_files_only": false,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 128000,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "TokenizersBackend"
+}