src/intent_model.py

"""
intent_model.py
----------------
TF-IDF + Logistic Regression intent classifier that routes free-text
warehouse queries into one of 8 operational intents. Chosen deliberately
over a heavier transformer classifier: it trains in <1s, needs no GPU/
internet on Spaces startup, and reaches high accuracy on this
template-generated-but-linguistically-varied dataset -- a good example of
picking the right-sized model for the job rather than defaulting to the
biggest one.
"""

from dataclasses import dataclass

import joblib
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline


@dataclass
class IntentPrediction:
    intent: str
    confidence: float


INTENT_DESCRIPTIONS = {
    "inventory_check": "Inventory / stock level lookup",
    "order_status": "Order status / tracking",
    "equipment_maintenance": "Equipment fault / maintenance request",
    "agv_navigation": "AGV / AMR routing & navigation",
    "picking_optimization": "Picking route / strategy optimization",
    "safety_incident": "Safety incident reporting",
    "system_status": "Equipment / system status check",
    "general_faq": "General warehouse automation question",
}


def build_pipeline() -> Pipeline:
    return Pipeline([
        ("tfidf", TfidfVectorizer(ngram_range=(1, 2), min_df=1, stop_words="english")),
        ("clf", LogisticRegression(max_iter=1000, C=8.0)),
    ])


def predict(pipeline: Pipeline, text: str) -> IntentPrediction:
    probs = pipeline.predict_proba([text])[0]
    classes = pipeline.classes_
    best_idx = probs.argmax()
    return IntentPrediction(intent=classes[best_idx], confidence=float(probs[best_idx]))


def load_pipeline(path: str) -> Pipeline:
    return joblib.load(path)


def save_pipeline(pipeline: Pipeline, path: str):
    joblib.dump(pipeline, path)