ml_model.py

"""
ML Pipeline Wrapper

Loads the fine-tuned mBERT archive from models/mbert_finetuned.keras and uses
it for sentiment predictions.
"""

import os
import json
import shutil
import tempfile
import zipfile
import re

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer


MODELS_DIR = "models"
CHECKPOINT_NAME = "mbert_finetuned.keras"
EXTRACTED_MODEL_DIR = os.path.join(tempfile.gettempdir(), "aims_nsight_mbert_finetuned")
EXTRACTED_MODEL_MARKER = os.path.join(EXTRACTED_MODEL_DIR, ".archive_signature")
MAX_LENGTH = 256

SENTIMENT_LABELS = {
    0: "Negative",
    1: "Neutral",
    2: "Positive",
}

TOPIC_LABELS = [
    "Course",
    "Internet",
    "Food",
    "Accommodation",
    "Library",
    "Sports",
    "Transportation",
    "Administration",
    "Health",
    "Security",
    "Events",
    "General",
    "Other",
]

TOPIC_KEYWORDS = {
    "Internet": [
        "internet", "wifi", "wi-fi", "network", "connection", "connexion",
        "connectivity", "bandwidth", "router", "online", "signal",
    ],
    "Food": [
        "food", "meal", "meals", "restaurant", "cafeteria", "canteen",
        "dining", "breakfast", "lunch", "dinner", "menu", "repas", "nourriture",
    ],
    "Accommodation": [
        "accommodation", "hostel", "room", "rooms", "dorm", "dormitory",
        "bed", "bathroom", "shower", "housing", "logement", "residence",
    ],
    "Library": [
        "library", "books", "book", "reading", "study room", "librarian",
        "bibliotheque", "bibliothèque",
    ],
    "Sports": [
        "sport", "sports", "football", "basketball", "gym", "fitness",
        "exercise", "training field", "terrain",
    ],
    "Transportation": [
        "transport", "transportation", "bus", "shuttle", "taxi", "car",
        "driver", "pickup", "drop-off", "travel", "commute",
    ],
    "Administration": [
        "administration", "admin", "registration", "registrar", "office",
        "document", "documents", "certificate", "payment", "fees", "finance",
        "visa", "scholarship", "staff",
    ],
    "Health": [
        "health", "clinic", "doctor", "nurse", "medical", "medicine",
        "hospital", "sick", "illness", "mental health", "counseling", "sante", "santé",
    ],
    "Security": [
        "security", "safe", "safety", "guard", "guards", "theft", "stolen",
        "harassment", "danger", "unsafe", "secure",
    ],
    "Events": [
        "event", "events", "seminar", "workshop", "conference", "ceremony",
        "party", "activity", "activities", "orientation",
    ],
    "Course": [
        "course", "courses", "class", "classes", "lecture", "lectures",
        "professor", "teacher", "lecturer", "assignment", "assignments",
        "exam", "exams", "grade", "grades", "curriculum", "module", "modules",
        "academic", "math", "mathematics", "statistics", "python", "project",
        "homework", "cours", "professeur",
    ],
}


class MLModel:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = None
        self.tokenizer = None
        self._load()

    def _load(self):
        path = os.path.join(MODELS_DIR, CHECKPOINT_NAME)
        if not os.path.exists(path):
            raise FileNotFoundError(f"Model archive not found: {path}")

        self._extract_model(path)
        self.tokenizer = AutoTokenizer.from_pretrained(
            EXTRACTED_MODEL_DIR,
            local_files_only=True,
            use_fast=True,
        )
        self.model = AutoModelForSequenceClassification.from_pretrained(
            EXTRACTED_MODEL_DIR,
            local_files_only=True,
        )

        if self.model.config.num_labels != len(SENTIMENT_LABELS):
            raise ValueError(
                f"Model has {self.model.config.num_labels} sentiment classes, but "
                f"{len(SENTIMENT_LABELS)} labels are configured"
            )

        self.model.to(self.device)
        self.model.eval()

        print(f"[ML] Fine-tuned mBERT model loaded from {path}")

    def _extract_model(self, archive_path):
        expected_files = ("config.json", "model.safetensors", "tokenizer.json")
        stat = os.stat(archive_path)
        signature = f"{os.path.abspath(archive_path)}:{stat.st_size}:{stat.st_mtime_ns}"

        if (
            all(os.path.exists(os.path.join(EXTRACTED_MODEL_DIR, name)) for name in expected_files)
            and os.path.exists(EXTRACTED_MODEL_MARKER)
            and self._read_archive_signature() == signature
        ):
            return

        if os.path.exists(EXTRACTED_MODEL_DIR):
            shutil.rmtree(EXTRACTED_MODEL_DIR)
        os.makedirs(EXTRACTED_MODEL_DIR, exist_ok=True)

        with zipfile.ZipFile(archive_path) as archive:
            archive.extractall(EXTRACTED_MODEL_DIR)

        self._write_label_config()
        with open(EXTRACTED_MODEL_MARKER, "w", encoding="utf-8") as marker:
            marker.write(signature)

    def _write_label_config(self):
        config_path = os.path.join(EXTRACTED_MODEL_DIR, "config.json")
        with open(config_path, encoding="utf-8") as config_file:
            config = json.load(config_file)

        config["id2label"] = {str(index): label for index, label in SENTIMENT_LABELS.items()}
        config["label2id"] = {label: index for index, label in SENTIMENT_LABELS.items()}

        with open(config_path, "w", encoding="utf-8") as config_file:
            json.dump(config, config_file, indent=2)
            config_file.write("\n")

    def _read_archive_signature(self):
        with open(EXTRACTED_MODEL_MARKER, encoding="utf-8") as marker:
            return marker.read()

    def predict(self, text: str) -> dict:
        if not self.model or not self.tokenizer:
            raise RuntimeError("ML model is not loaded")

        cleaned = (text or "").strip()
        if not cleaned:
            raise ValueError("Text is required")

        encoded = self.tokenizer(
            cleaned,
            return_tensors="pt",
            truncation=True,
            padding=True,
            max_length=MAX_LENGTH,
        )
        encoded = {key: value.to(self.device) for key, value in encoded.items()}

        with torch.no_grad():
            outputs = self.model(
                input_ids=encoded["input_ids"],
                attention_mask=encoded["attention_mask"],
            )

        sentiment_index = int(torch.argmax(outputs.logits, dim=1).item())

        try:
            sentiment_label = SENTIMENT_LABELS[sentiment_index]
        except KeyError as exc:
            raise ValueError(f"Model returned unknown sentiment class: {sentiment_index}") from exc
        sentiment_score = (
            1 if sentiment_label == "Positive"
            else (-1 if sentiment_label == "Negative" else 0)
        )

        return {
            "sentiment_class": sentiment_index,
            "sentiment_label": sentiment_label,
            "sentiment_score": sentiment_score,
            "topic": self._predict_topic(cleaned),
        }

    def _predict_topic(self, text: str) -> str:
        normalized = self._normalize_topic_text(text)
        scores = {}

        for topic, keywords in TOPIC_KEYWORDS.items():
            score = 0
            for keyword in keywords:
                normalized_keyword = self._normalize_topic_text(keyword)
                if " " in normalized_keyword:
                    if normalized_keyword in normalized:
                        score += len(normalized_keyword.split()) + 1
                elif re.search(rf"\b{re.escape(normalized_keyword)}\b", normalized):
                    score += 1

            if score:
                scores[topic] = score

        if not scores:
            return "General"

        return max(scores, key=lambda topic: (scores[topic], -TOPIC_LABELS.index(topic)))

    @staticmethod
    def _normalize_topic_text(text: str) -> str:
        text = text.lower()
        text = text.replace("é", "e").replace("è", "e").replace("ê", "e")
        text = text.replace("à", "a").replace("â", "a")
        text = text.replace("î", "i").replace("ï", "i")
        text = text.replace("ô", "o")
        text = text.replace("ù", "u").replace("û", "u")
        text = re.sub(r"[^a-z0-9\s-]", " ", text)
        return re.sub(r"\s+", " ", text).strip()