app.py

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

MODEL_ID = "gorkem371/pii-intent-classifier-xlmr-large"

print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
model.eval()
print("Model loaded!")

# ── Auto-detect entities from text ──────────────────────────────

ENTITY_PATTERNS = [
    ("EMAIL", r'[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}'),
    ("IBAN", r'\b[A-Z]{2}\d{2}[A-Z0-9]{10,30}\b'),
    ("PHONE", r'(?:\+?\d{1,3}[\s-]?)?\(?\d{2,4}\)?[\s.-]?\d{3,4}[\s.-]?\d{3,4}\b'),
    ("SOCIAL_MEDIA", r'@[a-zA-Z0-9_.]{2,30}'),
    ("URL", r'(?:https?://)?(?:www\.)?[a-zA-Z0-9-]+\.[a-zA-Z]{2,}(?:/\S*)?'),
    ("CRYPTO_ADDRESS", r'\b(?:0x[a-fA-F0-9]{40}|[13][a-km-zA-HJ-NP-Z1-9]{25,34}|bc1[a-zA-HJ-NP-Z0-9]{39,59})\b'),
]

INTENT_KEYWORDS = {
    "PHONE": [
        "numaram", "numara", "whatsapp", "telefon", "ara beni", "bana ulaş",
        "رقمي", "واتساب", "اتصل", "كلمني",
        "my number", "call me", "whatsapp me", "reach me", "contact me",
        "send you my number", "numara atarim", "numara veririm",
        "بارسل لك رقمي", "بعطيك رقمي",
    ],
    "SOCIAL_MEDIA": [
        "instagram", "insta", "tiktok", "telegram", "snapchat", "snap",
        "twitter", "takip et", "follow me", "تابعني", "انستقرام", "سناب",
        "beni bul", "find me", "تلقاني", "add me",
    ],
    "EMAIL": ["mail", "email", "ايميل", "راسلني", "e-posta"],
    "OFF_PLATFORM_ATTEMPT": [
        "başka yerde", "platform dışı", "dışarıda", "oradan konuşalım",
        "let's talk on", "move to", "نروح", "بره", "خارج",
        "bio", "story", "profil", "بايو", "ستوري",
    ],
}


def detect_entities(text):
    """Auto-detect entities in text."""
    found = []

    for etype, pattern in ENTITY_PATTERNS:
        for match in re.finditer(pattern, text):
            found.append((match.group(), etype, match.start()))

    # Remove overlapping matches (keep longest)
    found.sort(key=lambda x: (-len(x[0]), x[2]))
    used_ranges = []
    filtered = []
    for entity, etype, start in found:
        end = start + len(entity)
        if not any(s <= start < e or s < end <= e for s, e in used_ranges):
            filtered.append((entity, etype))
            used_ranges.append((start, end))

    return filtered


def detect_intent_type(text):
    """If no entity found, detect intent type from keywords."""
    text_lower = text.lower()
    for etype, keywords in INTENT_KEYWORDS.items():
        if any(kw in text_lower for kw in keywords):
            return etype
    return "PHONE"


def predict_single(context, entity, entity_type):
    """Run model on a single entity."""
    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)

    pii_prob = probs[0][1].item()
    not_pii_prob = probs[0][0].item()
    return pii_prob, not_pii_prob


def classify(message: str):
    if not message.strip():
        return {}, "Please enter a message."

    entities = detect_entities(message)

    if not entities:
        # No explicit entity — check for intent
        intent_type = detect_intent_type(message)
        entities = [("NONE", intent_type)]

    # Run model on all detected entities, take the highest PII score
    best_pii = 0
    best_not_pii = 1
    best_entity = ""
    best_type = ""
    all_results = []

    for entity, etype in entities:
        pii_prob, not_pii_prob = predict_single(message, entity, etype)
        all_results.append((entity, etype, pii_prob))
        if pii_prob > best_pii:
            best_pii = pii_prob
            best_not_pii = not_pii_prob
            best_entity = entity
            best_type = etype

    label_scores = {
        "🚨 PII DETECTED": best_pii,
        "✅ NOT PII": best_not_pii,
    }

    # Build verdict
    if best_pii > best_not_pii:
        verdict = f"🚨 **PII DETECTED** — This message contains PII sharing intent.\n\n"
    else:
        verdict = f"✅ **NOT PII** — No PII sharing intent detected.\n\n"

    verdict += f"**Confidence:** {max(best_pii, best_not_pii):.1%}\n\n"

    # Show what was detected
    if len(all_results) == 1 and all_results[0][0] == "NONE":
        verdict += f"📋 No explicit entity found — analyzed message intent as `{best_type}`"
    else:
        verdict += "📋 **Detected entities:**\n"
        for ent, etype, prob in all_results:
            icon = "🚨" if prob > 0.5 else "✅"
            verdict += f"- {icon} `{ent}` ({etype}) → {'PII' if prob > 0.5 else 'NOT PII'} ({prob:.0%})\n"

    return label_scores, verdict


EXAMPLES = [
    # PII - Direct sharing
    ["my number is 05321234567 call me anytime"],
    ["numaram 05321234567 bana yazabilirsin"],
    ["رقمي 0501234567 تواصل معي"],
    # PII - Social media / email
    ["follow me on instagram @creator_official"],
    ["email me at creator@gmail.com"],
    # PII - Implicit intent
    ["i will send you my whatsapp tomorrow"],
    ["check my bio my number is right there"],
    ["if we agree on terms i will share my number"],
    # PII - Humor + real number (V11 improvement)
    ["numaram pizza siparişi gibi 05321234567 haha 😂"],
    # PII - Old number still visible
    ["my old number was 05321234567 it might still work"],
    # NOT PII - Order/tracking numbers
    ["your order number is ORD-784321"],
    ["siparis numaraniz SPN-78432156"],
    # NOT PII - Room number / postal code (V11 improvement)
    ["oda numaram 532 otelde buluşalım"],
    ["posta kodu 34720 Kadıköy İstanbul"],
    # NOT PII - Scam warning / non-contact
    ["05321234567 this number is a scammer do not call"],
    ["bake at 180 degrees celsius for 45 minutes"],
    # NOT PII - Sarcasm
    ["my number is 00000000000 lol just kidding"],
]

with gr.Blocks(
    title="PII Intent Classifier",
    theme=gr.themes.Soft(),
) as demo:
    gr.Markdown(
        """
        # 🔍 PII Intent Classifier (V11)
        ### Multilingual context-aware PII detection (Turkish, Arabic, English)

        Type any message below. The model will automatically detect entities (phone numbers, emails, handles, etc.)
        and determine if the message contains **intent to share personal contact information**.

        **Model:** [gorkem371/pii-intent-classifier-xlmr-large](https://huggingface.co/gorkem371/pii-intent-classifier-xlmr-large) — XLM-RoBERTa Large, 550M params | **F1: 99.3%** · **Conversation accuracy: 97.0%** · Trained on 41K samples
        """
    )

    with gr.Row():
        with gr.Column(scale=2):
            message_input = gr.Textbox(
                label="Message",
                placeholder="Type a message to analyze... (Turkish, Arabic, or English)",
                lines=3,
            )
            submit_btn = gr.Button("🔍 Analyze", variant="primary", size="lg")

        with gr.Column(scale=1):
            label_output = gr.Label(label="Classification", num_top_classes=2)
            verdict_output = gr.Markdown(label="Details")

    submit_btn.click(
        fn=classify,
        inputs=[message_input],
        outputs=[label_output, verdict_output],
    )

    message_input.submit(
        fn=classify,
        inputs=[message_input],
        outputs=[label_output, verdict_output],
    )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[message_input],
        outputs=[label_output, verdict_output],
        fn=classify,
        cache_examples=False,
        label="💡 Try these examples",
    )

    gr.Markdown(
        """
        ---
        **How it works:** The model automatically detects entities (phone numbers, emails, @handles, URLs, IBANs) in your message
        using regex, then classifies each one for PII sharing intent using a fine-tuned XLM-RoBERTa Large model.
        If no entity is found, it analyzes the message for implicit intent (e.g., "I'll send my number later").

        **V11 improvements:** Understands humor + real numbers as PII, distinguishes room/postal/tracking numbers from real phone numbers, handles scam warnings with visible numbers, and 6 new NOT-PII categories.

        **Author:** [Gorkem Yildiz](https://gorkemyildiz.com) · **Dataset:** [pii-intent-detection-multilingual](https://huggingface.co/datasets/gorkem371/pii-intent-detection-multilingual) (41,427 samples)
        """
    )

demo.launch()