pii-intent-classifier-xlmr-large

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

model_name = "gorkem371/pii-intent-classifier-xlmr-large"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.eval()

def classify_pii(context: str, entity: str, entity_type: str) -> dict:
    """
    Classify whether a message contains PII sharing intent.

    Args:
        context: The full message text
        entity: The specific entity to classify (e.g., phone number, "NONE" if implicit)
        entity_type: Type of entity (PHONE, EMAIL, SOCIAL_MEDIA, IBAN, ADDRESS, URL, etc.)

    Returns:
        dict with 'is_pii' (bool) and 'confidence' (float)
    """
    text = f"{context} </s> {entity} | {entity_type}"
    inputs = tokenizer(text, max_length=256, padding="max_length", truncation=True, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**inputs)
        probs = F.softmax(outputs.logits, dim=-1)
        pred = torch.argmax(probs, dim=-1).item()
        confidence = probs[0][pred].item()

    return {
        "is_pii": pred == 1,
        "label": "PII" if pred == 1 else "NOT_PII",
        "confidence": round(confidence, 4)
    }

# Examples
print(classify_pii("my number is 05321234567 call me", "05321234567", "PHONE"))
# {'is_pii': True, 'label': 'PII', 'confidence': 0.9987}

print(classify_pii("order number is ORD-784321", "ORD-784321", "PHONE"))
# {'is_pii': False, 'label': 'NOT_PII', 'confidence': 0.9954}

print(classify_pii("i will send you my whatsapp tomorrow", "NONE", "PHONE"))
# {'is_pii': True, 'label': 'PII', 'confidence': 0.9821}

print(classify_pii("oda numaram 532 otelde buluşalım", "NONE", "PHONE"))
# {'is_pii': False, 'label': 'NOT_PII', 'confidence': 0.9876}