binga/privacy-filter-multitask

Privacy Filter Multi-Task 🔒📄

A single model for simultaneous PII Detection (NER) and Document Classification (10 categories).

Adapted from openai/privacy-filter — a 1.4B Sparse MoE transformer with only ~50M active parameters per token.

Architecture

Input → BPE Tokenizer (o200k_base, 200K vocab)
  ↓
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

PII Detection (NER)

8 entity types: private_person · private_email · private_phone · private_address · private_date · private_url · account_number · secret

Document Classification (10 classes)

Per-class test accuracy:

🚀 Production Inference Guide

All numbers below are measured on real hardware with both task heads (NER + doc classification) executing on every call. Benchmark script: single forward pass produces PII entity tags and document category simultaneously.

Resource Requirements

GPU — Single-Document Latency (NVIDIA A10G, bf16)

Time from raw text to both NER tags + document category:

Latency is dominated by a fixed ~105 ms kernel-launch overhead from the Sparse MoE routing — it barely changes with sequence length up to 512 tokens.

GPU — Batched Throughput (NVIDIA A10G, bf16)

GPU — Batched Latency Detail (NVIDIA A10G, bf16)

<details> <summary>Full latency table (click to expand)</summary>

</details>

GPU — Peak VRAM Usage (bf16)

The model is extremely memory-efficient. Even at batch=64, seq=512, it uses only 6.2 GB — comfortably fits on a T4 (16 GB). This is because the Sparse MoE only activates 4 of 128 experts per token.

CPU — Latency & Throughput (AMD EPYC 7R32, 8 cores, fp32)

On CPU the model runs at ~152 ms/doc for short texts (seq=64, bs=1) — suitable for low-volume or batch-offline pipelines.

Daily Throughput Projections

Sustained throughput for a single device, running 24/7 at the optimal batch size:

Multi-GPU Scaling Estimates

_{¹ A100 estimates are linearly extrapolated from A10G numbers using A100's ~2.3× higher memory bandwidth and larger batch capacity. Actual numbers will vary — benchmark on your target hardware.}

Serving Recommendations

_{² At seq=64 most documents will be truncated. Use seq=128–256 for production balance.}

Key observations:

• —The model has a fixed ~105 ms overhead per forward pass regardless of sequence length (MoE routing + expert dispatch). Batching amortizes this cost across documents — the per-doc cost drops from 106 ms (bs=1) to under 10 ms (bs=64).
• —Memory is not the bottleneck — even at bs=64/seq=512 the model uses only 6.2 GB. You can run this on a T4 (16 GB) with room to spare.
• —Optimal batch size for throughput: bs=64 for all sequence lengths on A10G.
• —Optimal batch size for latency-constrained: bs=8–16 gives a good per-doc latency (13–19 ms) while keeping batch latency under 225 ms.

Training Strategy

Two-phase training approach:

1. 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.

1. 1.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:
2. 2.Preserves NER performance exactly (backbone untouched)
3. 3.Is extremely fast (~seconds per epoch on cached features)
4. 4.Achieves 47.8% test accuracy (up from 24.8% in phase 1)

Usage

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()

# 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():
    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%})")

Batched Inference (Production)

# Process a batch of documents — both tasks in a single forward pass
texts = ["doc1...", "doc2...", "doc3...", ...]
inputs = tokenizer(texts, return_tensors="pt", padding=True,
                   truncation=True, max_length=256).to(model.device)

with torch.no_grad():
    outputs = model(**inputs, output_hidden_states=True)

# NER predictions for all docs: [batch, seq_len]
ner_preds = outputs.logits.argmax(dim=-1)

# Doc class for all docs: [batch]
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)
doc_preds = doc_head(pooled).argmax(dim=-1)

Example Outputs

Files