inference.py

"""Vertical + Sub-vertical classifier — inference module.
Loads once, exposes classify(category_name) -> dict.
Vertical: closed set of 21 (never 'others'); low confidence -> review flag.
Sub-vertical: masked to the predicted vertical's children; low confidence -> propose-new flag.
"""
import re, json, joblib, numpy as np
from scipy.sparse import hstack

VERT_THR = 0.65   # below -> route vertical to human review
SUB_THR  = 0.50   # below -> propose a NEW sub-vertical

_V  = joblib.load("vertical_model.joblib")     # word_vec, char_vec, clf, classes
_S  = joblib.load("subvertical_model.joblib")  # clf (LinearSVC), classes
_P  = json.load(open("parent_map.json"))       # sub_vertical -> vertical
_WV, _CV, _VCLF = _V["word_vec"], _V["char_vec"], _V["clf"]
_VCLASSES = np.array(_VCLF.classes_)
_SCLASSES = np.array(_S["classes"])
_SPARENT  = np.array([_P.get(c, "") for c in _SCLASSES])

def _clean(s): return " ".join(str(s).lower().split())
def _vec(s):   return hstack([_WV.transform([s]), _CV.transform([s])]).tocsr()

def _suggest_name(name: str) -> str:
    """Derive a candidate sub-vertical name from the category text (offline heuristic)."""
    s = str(name)
    s = re.sub(r"\([^)]*\)", "", s)                       # drop "(Malayalam Movie)" etc.
    s = re.sub(r"\s+[Ii]n\s+[A-Z][a-zA-Z]+(\s+[A-Z][a-zA-Z]+)*\s*$", "", s)  # trailing "In <City>"
    if "-" in s:                                           # "Product-Brand" -> keep product
        s = s.split("-")[0]
    s = " ".join(s.lower().split()).strip(" -/&,.")
    return s or _clean(name)

def classify(name: str, vert_thr=VERT_THR, sub_thr=SUB_THR) -> dict:
    x = _vec(_clean(name))
    # ---- vertical ----
    vp = _VCLF.predict_proba(x)[0]
    order = vp.argsort()[::-1]
    vertical = _VCLASSES[order[0]]; v_conf = float(vp[order[0]])
    top3 = [(_VCLASSES[i], round(float(vp[i]), 4)) for i in order[:3]]
    # ---- sub-vertical (masked to this vertical's children) ----
    mask = _SPARENT == vertical
    if mask.any():
        m = _S["clf"].decision_function(x)[0].copy()
        m[~mask] = -1e9
        e = np.exp(m - m.max()); e[~mask] = 0.0; p = e / e.sum()
        j = int(p.argmax()); sub = _SCLASSES[j]; s_conf = float(p[j])
    else:
        sub, s_conf = None, 0.0
    propose_new = s_conf < sub_thr
    return {
        "input": name,
        "vertical": vertical,
        "vertical_confidence": round(v_conf, 4),
        "vertical_review": v_conf < vert_thr,
        "vertical_top3": top3,
        "sub_vertical": sub,
        "sub_vertical_confidence": round(s_conf, 4),
        "propose_new_sub_vertical": propose_new,
        "suggested_sub_vertical": _suggest_name(name) if propose_new else None,
    }

if __name__ == "__main__":
    tests = ["Cardiologist Doctors", "Bridal Makeup Artists", "iPhone Repair Service",
             "Gold Necklace Dealers-Tanishq", "Yoga Classes For Kids",
             "Pakal Pole (Malayalam Movie)", "Drone Camera Dealers-DJI",
             "Cattle Feed Suppliers", "Sourdough Bread Bakery", "Wedding Photographer"]
    for t in tests:
        r = classify(t)
        flag = "  [REVIEW]" if r["vertical_review"] else ""
        newf = "  [PROPOSE-NEW-SUB]" if r["propose_new_sub_vertical"] else ""
        print(f'{t:38s} -> {r["vertical"]:22s} {r["vertical_confidence"]:.2f}{flag} | '
              f'{str(r["sub_vertical"]):32s} {r["sub_vertical_confidence"]:.2f}{newf}')