predictor.py

"""ScamShield NLP — inference-only predictor.

Loads the already-trained calibrated Linear SVM and the fitted TF-IDF
vectorizer from ``models/`` and exposes a single ``predict(text)`` method.

No training code, no dataset access, no re-fitting — the Space is a pure
inference service. Both artifacts come directly from the original experiment
(``models/svm.joblib`` and ``models/tfidf_vectorizer.joblib``).

Model details (verified against the original training script,
``src/models/train_nlp_baselines.py``):

* Classifier : ``CalibratedClassifierCV(LinearSVC(C=1.0))``
  (``predict_proba`` returns genuine sigmoid-calibrated probabilities — the
  confidence value is NOT fabricated).
* Class order: ``[0, 1]`` — class 1 is the scam/phishing class.
* Vectorizer: ``TfidfVectorizer(ngram_range=(1, 2), max_features=10000,
  min_df=2, sublinear_tf=True, lowercase=True)``.
"""

import os

import joblib

from preprocessing import clean_text

MODELS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
VECTORIZER_PATH = os.path.join(MODELS_DIR, "tfidf_vectorizer.joblib")
MODEL_PATH = os.path.join(MODELS_DIR, "svm.joblib")

# Label threshold used in the original predictor (probability of class 1).
THRESHOLD = 0.5


class ScamPredictor:
    """Inference wrapper around the trained TF-IDF + calibrated Linear SVM."""

    def __init__(self):
        self.vectorizer = joblib.load(VECTORIZER_PATH)
        self.model = joblib.load(MODEL_PATH)

    @staticmethod
    def risk_level(prob: float) -> str:
        """Map scam probability to a risk band (same logic as the experiment)."""
        if prob >= 0.8:
            return "HIGH"
        if prob >= 0.5:
            return "MEDIUM"
        return "LOW"

    def predict(self, text: str) -> dict:
        """Classify a single message and return prediction, confidence, risk.

        Returns:
            {
                "prediction": "SCAM" | "LEGITIMATE",
                "confidence": <float 0..1>,   # calibrated P(scam)
                "risk_level": "HIGH" | "MEDIUM" | "LOW",
            }
        """
        cleaned = clean_text(text)
        X = self.vectorizer.transform([cleaned])
        prob_scam = float(self.model.predict_proba(X)[0, 1])
        prediction = "SCAM" if prob_scam >= THRESHOLD else "LEGITIMATE"
        return {
            "prediction": prediction,
            "confidence": round(prob_scam, 4),
            "risk_level": self.risk_level(prob_scam),
        }