src/predict.py

from __future__ import annotations

from pathlib import Path

import joblib
from sentence_transformers import SentenceTransformer


DEFAULT_MODEL_DIR = Path("models")


def load_pipeline(model_dir: Path = DEFAULT_MODEL_DIR):
    """Load the saved embedding model name and trained classifier."""
    embedding_model_name = (model_dir / "embedding_model.txt").read_text(encoding="utf-8").strip()
    embedder = SentenceTransformer(embedding_model_name)
    classifier = joblib.load(model_dir / "classifier.joblib")
    return embedder, classifier


def predict(text: str, model_dir: Path = DEFAULT_MODEL_DIR) -> tuple[str, dict[str, float]]:
    """Predict one job posting and return the label plus confidence scores."""
    embedder, classifier = load_pipeline(model_dir)
    embedding = embedder.encode([text], normalize_embeddings=True)
    label = classifier.predict(embedding)[0]

    if hasattr(classifier, "predict_proba"):
        probabilities = classifier.predict_proba(embedding)[0]
        scores = {
            str(class_name): float(probability)
            for class_name, probability in zip(classifier.classes_, probabilities)
        }
    else:
        scores = {label: 1.0}

    return label, scores