| """Shared model loading and inference for the API, Gradio and other clients.""" |
|
|
| import os |
| from pathlib import Path |
|
|
| import torch |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
| PROJECT_ROOT = Path(__file__).resolve().parent.parent |
| DEFAULT_MODEL_ID = str(PROJECT_ROOT / "model" / "sentiment_model") |
| MAX_LENGTH = 256 |
|
|
|
|
| class SentimentAnalyzer: |
| """Load the fine-tuned model once and expose a reusable prediction method.""" |
|
|
| def __init__(self, model_id: str | None = None, device: str | None = None): |
| self.model_id = model_id or os.getenv("MODEL_ID", DEFAULT_MODEL_ID) |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| self.tokenizer = None |
| self.model = None |
|
|
| @property |
| def is_loaded(self) -> bool: |
| return self.model is not None and self.tokenizer is not None |
|
|
| def load(self) -> "SentimentAnalyzer": |
| if self.is_loaded: |
| return self |
|
|
| model_path = Path(self.model_id) |
| if model_path.is_absolute() and not model_path.exists(): |
| raise FileNotFoundError( |
| f"Modele introuvable dans {model_path}. Lancez d'abord scripts/train.py." |
| ) |
|
|
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_id) |
| self.model = AutoModelForSequenceClassification.from_pretrained( |
| self.model_id |
| ).to(self.device) |
| self.model.eval() |
| return self |
|
|
| def unload(self) -> None: |
| self.model = None |
| self.tokenizer = None |
|
|
| def predict(self, text: str, include_debug: bool = False) -> dict: |
| if not self.is_loaded: |
| raise RuntimeError("Le modele doit etre charge avant la prediction.") |
|
|
| inputs = self.tokenizer( |
| text, |
| return_tensors="pt", |
| truncation=True, |
| max_length=MAX_LENGTH, |
| ).to(self.device) |
|
|
| with torch.no_grad(): |
| logits = self.model(**inputs).logits[0] |
|
|
| probabilities = torch.softmax(logits, dim=-1) |
| predicted_class_id = int(torch.argmax(probabilities)) |
|
|
| result = { |
| "label": self.model.config.id2label[predicted_class_id], |
| "confidence": round(float(probabilities[predicted_class_id]), 4), |
| "probabilities": { |
| self.model.config.id2label[i]: round(float(probability), 4) |
| for i, probability in enumerate(probabilities) |
| }, |
| } |
| if include_debug: |
| result["debug"] = { |
| "tokens": self.tokenizer.convert_ids_to_tokens( |
| inputs["input_ids"][0].detach().cpu().tolist() |
| ), |
| "logits": [round(float(logit), 4) for logit in logits], |
| } |
| return result |
|
|
