car_advisor/nlp_model.py

from typing import Dict, List
from .config import DEFAULT_LABELS

KEYWORD_MAP = {
    "brake_wear": ["brake", "squeal", "screech", "stopping", "pads"],
    "flat_tire": ["flat", "puncture", "tyre", "tire", "pressure", "air"],
    "engine_leak": ["leak", "oil", "puddle", "drip", "smell burning"],
    "cracked_windshield": ["crack", "windshield", "glass"],
    "paint_damage": ["scratch", "scrape", "paint", "scuff"],
    "scratch_dent": ["dent", "dented", "bent"],
    "headlight_fault": ["headlight", "bulb", "beam", "lamp"],
    "battery_corrosion": ["battery", "corrosion", "terminal", "start"],
    "rust": ["rust", "oxid"],
    "bumper_damage": ["bumper", "fender"]
}

def _contains_any(text: str, keywords: List[str]) -> bool:
    t = (text or "").lower()
    return any(kw in t for kw in keywords)

class NLPInference:
    def __init__(self, labels: List[str] = None, ckpt_dir: str = "checkpoints/nlp"):
        self.labels = labels or DEFAULT_LABELS
        self.ckpt_dir = ckpt_dir
        self.trained = False
        try:
            import joblib, os
            clf_p = os.path.join(ckpt_dir, "best.joblib")
            mlb_p = os.path.join(ckpt_dir, "mlb.joblib")
            if os.path.exists(clf_p) and os.path.exists(mlb_p):
                self.clf = joblib.load(clf_p)
                self.mlb = joblib.load(mlb_p)
                self.trained = True
        except Exception:
            self.trained = False

    def predict(self, text: str) -> Dict[str, float]:
        if not text:
            return {l: 0.0 for l in self.labels}
        if self.trained:
            probs = self.clf.predict_proba([text])[0]
            out = {}
            for i, lbl in enumerate(self.mlb.classes_):
                val = probs[i] if isinstance(probs[i], (float,int)) else probs[i][1]
                out[lbl] = float(val)
            for l in self.labels:
                out.setdefault(l, 0.0)
            return out
        else:
            scores = {l: 0.01 for l in self.labels}
            for label, kws in KEYWORD_MAP.items():
                if _contains_any(text, kws):
                    scores[label] += 0.5
            s = sum(scores.values())
            return {k: v/s for k,v in scores.items()}