src/inference.py

"""Shared inference helpers: PII redaction -> tokenize -> classify.

Used by both the training notebook (for the inference demo) and app.py (for the
deployed Gradio app), so there is no train/serve skew in how predictions are made.
"""

import torch
from sklearn.preprocessing import LabelEncoder
from transformers import PreTrainedModel, PreTrainedTokenizerBase

from src.pii import RedactionResult, redact_pii

MAX_LENGTH = 256


def predict_intent(
    text: str,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    label_encoder: LabelEncoder,
    device: torch.device,
    top_k: int = 3,
) -> list[tuple[str, float]]:
    """Top-k (intent, probability) predictions for already-redacted text."""
    model.eval()
    encoded = tokenizer(text, truncation=True, max_length=MAX_LENGTH, return_tensors="pt").to(device)
    with torch.no_grad():
        logits = model(**encoded).logits.squeeze(0)
    probs = torch.softmax(logits, dim=-1).cpu().numpy()
    top_idx = probs.argsort()[::-1][:top_k]
    return [(label_encoder.inverse_transform([i])[0], float(probs[i])) for i in top_idx]


def predict_safe(
    text: str,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    label_encoder: LabelEncoder,
    device: torch.device,
    top_k: int = 3,
    hash_identifiers: bool = True,
) -> tuple[RedactionResult, list[tuple[str, float]]]:
    """Redact PII, then classify the redacted text. Never sends raw PII to the model."""
    redaction = redact_pii(text, hash_identifiers=hash_identifiers)
    predictions = predict_intent(redaction.redacted_text, model, tokenizer, label_encoder, device, top_k)
    return redaction, predictions