inference.py

from __future__ import annotations

from pathlib import Path
from typing import Any

import joblib


MODEL_PATH = Path(__file__).resolve().parent / "model" / "language_classifier.joblib"
model = joblib.load(MODEL_PATH)


def _normalize_text(text: Any) -> str:
    return " ".join(str(text).replace("\ufeff", " ").split())


def _as_text_list(inputs: str | list[str]) -> list[str]:
    if isinstance(inputs, str):
        return [_normalize_text(inputs)]
    return [_normalize_text(text) for text in inputs]


def predict(inputs: str | list[str]) -> list[dict[str, float | str]]:
    texts = _as_text_list(inputs)
    if not texts:
        return []

    labels = [str(label) for label in model.predict(texts)]

    if hasattr(model, "predict_proba") and hasattr(model, "classes_"):
        classes = [str(label) for label in model.classes_]
        probabilities = model.predict_proba(texts)
        return [
            {
                "label": label,
                "score": round(float(row[classes.index(label)]), 6),
            }
            for label, row in zip(labels, probabilities)
        ]

    return [{"label": label, "score": 1.0} for label in labels]


def predict_all_scores(inputs: str | list[str]) -> list[dict[str, Any]]:
    texts = _as_text_list(inputs)
    if not texts:
        return []

    labels = [str(label) for label in model.predict(texts)]

    if not hasattr(model, "predict_proba") or not hasattr(model, "classes_"):
        return [
            {"label": label, "score": 1.0, "scores": {label: 1.0}}
            for label in labels
        ]

    classes = [str(label) for label in model.classes_]
    probabilities = model.predict_proba(texts)
    results = []
    for label, row in zip(labels, probabilities):
        scores = {
            class_label: round(float(score), 6)
            for class_label, score in zip(classes, row)
        }
        results.append(
            {
                "label": label,
                "score": scores[label],
                "scores": scores,
            }
        )
    return results