Spaces:
Paused
Paused
File size: 1,512 Bytes
f1f47b0 4b1c4d0 f1f47b0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | """
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)) |