Inference.py

"""
inference.py
Reusable inference module for spam classification.
Exposes a single predict(text) function so it can be imported
into any notebook, script, or the Gradio app without duplicating
model-loading logic.
"""

from transformers import pipeline

# Chosen model after benchmarking (see report.md for rationale).
# NOTE: update this if niru-nny/SMS_Spam_Detection scores higher once you
# re-run model_benchmark.ipynb with the corrected candidate list.
MODEL_NAME = "wesleyacheng/sms-spam-classification-with-bert"

_classifier = None


def _get_classifier():
    """Lazy-load the pipeline once and reuse it across calls."""
    global _classifier
    if _classifier is None:
        _classifier = pipeline("text-classification", model=MODEL_NAME)
    return _classifier


def predict(text: str) -> dict:
    """
    Run spam classification on a single piece of text.

    Args:
        text: the message to classify.

    Returns:
        dict with keys: 'label' ('spam' or 'ham') and 'score' (float, 0-1).
    """
    clf = _get_classifier()
    result = clf(text)[0]
    label = "spam" if result["label"] in ("LABEL_1", "spam", "SPAM") else "ham"
    return {"label": label, "score": round(result["score"], 4)}


if __name__ == "__main__":
    samples = [
        "Congratulations! You've won a $1000 gift card, click here to claim now!",
        "Hey, are we still on for lunch tomorrow?",
    ]
    for s in samples:
        print(s, "->", predict(s))