Update README with head-only doc clf results (test_acc=0.478)

README.md

@@ -19,108 +19,176 @@ model-index:
   results:
   - task:
       type: token-classification
-      name: PII Detection
+      name: PII Detection (NER)
     dataset:
       name: ai4privacy/pii-masking-400k
       type: ai4privacy/pii-masking-400k
     metrics:
     - type: f1
       value: 0.4925
+      name: F1 (strict span-level)
     - type: precision
       value: 0.6968
     - type: recall
       value: 0.3809
   - task:
       type: text-classification
-      name: Document Classification
+      name: Document Classification (10 classes)
     dataset:
       name: yahoo_answers_topics
       type: community-datasets/yahoo_answers_topics
     metrics:
     - type: accuracy
-      value: 0.2482
+      value: 0.4776
+      name: Test Accuracy
 ---
 
 # 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).
+Adapted from [openai/privacy-filter](https://huggingface.co/openai/privacy-filter) — a 1.4B Sparse MoE transformer with only ~50M active parameters per token.
 
 ## Architecture
 
 ```
-Input → BPE Tokenizer (200K vocab)
+Input → BPE Tokenizer (o200k_base, 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
+8-layer Sparse MoE Transformer
+  • 128 experts, top-4 routing (~50M active params/token)
+  • Banded sliding-window attention (window=128)
+  • GQA: 14 query heads, 2 KV heads, head_dim=64
+  • Hidden size: 640
+  ↓                          ↓
+NER Head (640→33)        Doc Head (mean-pool → 640→10)
+  ↓                          ↓
+BIOES PII tags            10-class document category
 ```
 
 ## 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** |
+### PII Detection (NER)
+
+| Metric | Value |
+|--------|-------|
+| **F1 (strict span-level)** | **0.493** |
+| Precision | 0.697 |
+| Recall | 0.381 |
+| Token Accuracy | 0.944 |
+
+8 entity types: `private_person` · `private_email` · `private_phone` · `private_address` · `private_date` · `private_url` · `account_number` · `secret`
+
+### Document Classification (10 classes)
+
+| Split | Accuracy |
+|-------|----------|
+| Val | 0.470 |
+| **Test** | **0.478** |
+
+Per-class test accuracy:
+
+| Category | Accuracy |
+|----------|----------|
+| Computers & Internet | 0.688 |
+| Family & Relationships | 0.615 |
+| Science & Mathematics | 0.556 |
+| Health | 0.524 |
+| Sports | 0.523 |
+| Politics & Government | 0.493 |
+| Entertainment & Music | 0.444 |
+| Society & Culture | 0.363 |
+| Education & Reference | 0.310 |
+| Business & Finance | 0.263 |
 
 ### Inference Speed
+
 | Device | Latency |
 |--------|---------|
-| GPU A10G (bf16) | 178 ms |
-| CPU (fp32) | 202 ms |
+| **GPU (A10G, bf16)** | **~154 ms/sample** |
+
+## Training Strategy
+
+Two-phase training approach:
 
-## PII Entity Types
-`private_person` • `private_email` • `private_phone` • `private_address` • `private_date` • `private_url` • `account_number` • `secret`
+1. **Phase 1 — Multi-task fine-tuning**: Partially unfroze last 4 MoE layers + both task heads. Trained on 20K NER examples (ai4privacy) + 20K doc examples (Yahoo Answers). Multi-task loss (NER×1.0 + Doc×0.5). 2 epochs, LR=2e-5.
 
-## Document Categories
-Society & Culture • Science & Math • Health • Education • Computers & Internet • Sports • Business & Finance • Entertainment • Family • Politics
+2. **Phase 2 — Doc head retraining** (head-only): Froze entire backbone + NER head. Pre-computed 640-dim pooled features for 100K Yahoo Answers examples. Trained fresh `Linear(640→10)` classifier for 10 epochs, LR=1e-3, cosine decay. This approach:
+   - Preserves NER performance exactly (backbone untouched)
+   - Is extremely fast (~seconds per epoch on cached features)
+   - Achieves **47.8% test accuracy** (up from 24.8% in phase 1)
 
 ## Usage
 
 ```python
-import torch, torch.nn as nn
+import torch
+import torch.nn as nn
 from transformers import AutoModelForTokenClassification, AutoTokenizer
 from huggingface_hub import hf_hub_download
 
+# Load model + tokenizer
 tokenizer = AutoTokenizer.from_pretrained("binga/privacy-filter-multitask")
 model = AutoModelForTokenClassification.from_pretrained(
     "binga/privacy-filter-multitask", dtype=torch.bfloat16, device_map="auto"
 )
+
+# Load document classification head
 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)
+doc_head.eval()
 
-text = "John Smith (SSN: 123-45-6789) emailed john@corp.com"
+# Inference
+text = "John Smith (SSN: 123-45-6789) emailed john@corp.com about Q3 earnings."
 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()]}")
+    outputs = model(**inputs, output_hidden_states=True)
+
+# === PII Detection ===
+print("PII entities:")
+for tok, pred in zip(
+    tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]),
+    outputs.logits.argmax(-1)[0]
+):
+    label = model.config.id2label[pred.item()]
+    if label != "O":
+        print(f"  {tok} → {label}")
+
+# === Document Classification ===
+categories = [
+    "Society & Culture", "Science & Math", "Health", "Education",
+    "Computers & Internet", "Sports", "Business & Finance",
+    "Entertainment", "Family", "Politics"
+]
+hidden = outputs.hidden_states[-1]
+mask = inputs["attention_mask"].unsqueeze(-1).to(hidden.dtype)
+pooled = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1)
+probs = torch.softmax(doc_head(pooled)[0].float(), dim=-1)
+top = probs.argmax().item()
+print(f"\nCategory: {categories[top]} ({probs[top]:.1%})")
 ```
 
-## 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
+## Example Outputs
+
+| Input | PII Detected | Category (confidence) |
+|-------|-------------|----------------------|
+| "My name is John Smith... email john@example.com" | ✅ John Smith, john@example.com, 123 Main St | Computers & Internet (56%) |
+| "Liverpool FC defeated Manchester City 3-1" | ❌ None | **Sports (98%)** |
+| "Federal Reserve announced a rate cut" | ❌ None | **Politics (52%)** |
+| "health benefits of meditation and yoga" | ❌ None | **Health (38%)** |
+| "Patient Jane Doe (SSN: 123-45-6789)" | ✅ Jane Doe, 123-45-6789, jane.doe@hospital.com | Education (41%) |
+| "learn programming? I want to learn Python" | ❌ None | **Education (53%)** |
+| "legal to record phone calls in California?" | ❌ None | **Politics (64%)** |
+
+## Files
+
+| File | Size | Description |
+|------|------|-------------|
+| `model.safetensors` | 2.6 GB | Backbone + NER head (1.4B MoE params) |
+| `doc_head.pt` | 26 KB | Document classification head (640→10) |
+| `config.json` | 3 KB | Model architecture config |
+| `tokenizer.json` | 27 MB | BPE tokenizer (o200k_base) |
+| `multitask_config.json` | 349 B | Multi-task metadata |