from peft import PeftModel
from transformers import AutoModelForTokenClassification, AutoTokenizer

# Load model
base_model = AutoModelForTokenClassification.from_pretrained(
    "jhu-clsp/mmBERT-base", num_labels=35  # O + 17 entity types × 2 (B/I)
)
model = PeftModel.from_pretrained(base_model, "llm-semantic-router/mmbert-pii-detector-lora")
tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/mmBERT-base")

# Detect PII
text = "Contact John Smith at john.smith@email.com or call 555-123-4567"
inputs = tokenizer(text.split(), is_split_into_words=True, return_tensors="pt")
outputs = model(**inputs)
predictions = outputs.logits.argmax(-1)

# Map predictions to labels
# ... (see full example in repository)

Supported PII Types

Entity Type	Description
PERSON	Full names
EMAIL_ADDRESS	Email addresses
PHONE_NUMBER	Phone numbers (various formats)
STREET_ADDRESS	Physical addresses
CREDIT_CARD	Credit card numbers
US_SSN	US Social Security Numbers
IP_ADDRESS	IP addresses
DATE_TIME	Dates and times
URL	Web URLs
ORGANIZATION	Company/org names

Use Cases

Privacy Compliance: GDPR, CCPA compliance
Data Masking: Redact PII before processing
LLM Safety: Prevent PII leakage in LLM outputs
Document Processing: Automated PII detection in documents

Part of vLLM Semantic Router

This model is part of the vLLM Semantic Router project.

License

Apache 2.0