Update training script and model files

services/label_model_manage.py

@@ -3,7 +3,6 @@
 import os
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from huggingface_hub import hf_hub_download
 import logging
 
 logger = logging.getLogger(__name__)
@@ -21,7 +20,7 @@ class KpaModelManager:
         self.model_id = None
 
     def load_model(self, model_id: str, api_key: str = None):
-        """Load model with weights from Hugging Face repository"""
+        """Load complete model and tokenizer directly from Hugging Face"""
         if self.model_loaded:
             logger.info("KPA model already loaded")
             return
@@ -39,35 +38,20 @@ class KpaModelManager:
             # Prepare token for authentication if API key is provided
             token = api_key if api_key else None
             
-            # Load base tokenizer (distilbert-base-uncased)
-            base_model_name = "distilbert-base-uncased"
-            logger.info(f"Loading tokenizer from {base_model_name}...")
-            self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-            
-            # Load base model architecture
-            logger.info(f"Loading base model architecture from {base_model_name}...")
-            self.model = AutoModelForSequenceClassification.from_pretrained(
-                base_model_name,
-                num_labels=2
+            # Load tokenizer and model directly from Hugging Face
+            logger.info("Loading tokenizer...")
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                model_id,
+                token=token,
+                trust_remote_code=True
             )
             
-            # Download and load fine-tuned weights from Hugging Face
-            logger.info(f"Downloading fine-tuned weights from {model_id}...")
-            weights_path = hf_hub_download(
-                repo_id=model_id,
-                filename="modele_appariement_rapide.pth",
-                token=token
+            logger.info("Loading model...")
+            self.model = AutoModelForSequenceClassification.from_pretrained(
+                model_id,
+                token=token,
+                trust_remote_code=True
             )
-            
-            logger.info(f"Loading fine-tuned weights from {weights_path}...")
-            checkpoint = torch.load(weights_path, map_location=self.device)
-            
-            # Load state dict
-            if "model_state_dict" in checkpoint:
-                self.model.load_state_dict(checkpoint["model_state_dict"])
-            else:
-                self.model.load_state_dict(checkpoint)
-            
             self.model.to(self.device)
             self.model.eval()
 
@@ -118,6 +102,10 @@ class KpaModelManager:
             raise RuntimeError(f"KPA prediction failed: {str(e)}")
         
     def get_model_info(self):
+        """Get model information"""
+        if not self.model_loaded:
+            return {"loaded": False}
+            
         return {
             "model_name": self.model_id,
             "device": str(self.device),
@@ -127,5 +115,5 @@ class KpaModelManager:
         }
 
 
-
-kpa_model_manager = KpaModelManager()
+# Initialize singleton instance
+kpa_model_manager = KpaModelManager()

services/stance_model_manager.py

@@ -1,4 +1,4 @@
-"""Model manager for keypoint–argument matching model"""
+"""Model manager for stance detection model"""
 
 import os
 import torch
@@ -8,39 +8,34 @@ import logging
 logger = logging.getLogger(__name__)
 
 
-class KpaModelManager:
-    """Manages loading and inference for keypoint matching model"""
-
+class StanceModelManager:
+    """Manages stance detection model loading and predictions"""
+    
     def __init__(self):
         self.model = None
         self.tokenizer = None
         self.device = None
         self.model_loaded = False
-        self.max_length = 256
-        self.model_id = None
-
+    
     def load_model(self, model_id: str, api_key: str = None):
-        """Load complete model and tokenizer directly from Hugging Face"""
+        """Load model and tokenizer from Hugging Face"""
         if self.model_loaded:
-            logger.info("KPA model already loaded")
+            logger.info("Stance model already loaded")
             return
-
+        
         try:
-            logger.info(f"Loading KPA model from Hugging Face: {model_id}")
+            logger.info(f"Loading stance model from Hugging Face: {model_id}")
             
             # Determine device
             self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
             logger.info(f"Using device: {self.device}")
             
-            # Store model ID
-            self.model_id = model_id
-            
             # Prepare token for authentication if API key is provided
             token = api_key if api_key else None
             
-            # Load tokenizer and model directly from Hugging Face
+            # Load tokenizer and model from Hugging Face
             logger.info("Loading tokenizer...")
-            self.tokenizer = AutoTokenizer.from_pretrained(
+            self.tokenizer = AutoTokenizer.from_pretrained( 
                 model_id,
                 token=token,
                 trust_remote_code=True
@@ -54,66 +49,53 @@ class KpaModelManager:
             )
             self.model.to(self.device)
             self.model.eval()
-
+            
             self.model_loaded = True
-            logger.info("✓ KPA model loaded successfully from Hugging Face!")
-
+            logger.info("✓ Stance model loaded successfully from Hugging Face!")
+            
         except Exception as e:
-            logger.error(f"Error loading KPA model: {str(e)}")
-            raise RuntimeError(f"Failed to load KPA model: {str(e)}")
-
-    def predict(self, argument: str, key_point: str) -> dict:
-        """Run a prediction for (argument, key_point)"""
+            logger.error(f"Error loading stance model: {str(e)}")
+            raise RuntimeError(f"Failed to load stance model: {str(e)}")
+    
+    def predict(self, topic: str, argument: str) -> dict:
+        """Make a single stance prediction"""
         if not self.model_loaded:
-            raise RuntimeError("KPA model not loaded")
-
-        try:
-            # Tokenize input
-            encoding = self.tokenizer(
-                argument,
-                key_point,
-                truncation=True,
-                padding="max_length",
-                max_length=self.max_length,
-                return_tensors="pt"
-            ).to(self.device)
-
-            # Forward pass
-            with torch.no_grad():
-                outputs = self.model(**encoding)
-                logits = outputs.logits
-                probabilities = torch.softmax(logits, dim=-1)
-
-                predicted_class = torch.argmax(probabilities, dim=-1).item()
-                confidence = probabilities[0][predicted_class].item()
-
-            return {
-                "prediction": predicted_class,
-                "confidence": confidence,
-                "label": "apparie" if predicted_class == 1 else "non_apparie",
-                "probabilities": {
-                    "non_apparie": probabilities[0][0].item(),
-                    "apparie": probabilities[0][1].item(),
-                },
-            }
-
-        except Exception as e:
-            logger.error(f"Error during prediction: {str(e)}")
-            raise RuntimeError(f"KPA prediction failed: {str(e)}")
+            raise RuntimeError("Stance model not loaded")
+        
+        # Format input
+        text = f"Topic: {topic} [SEP] Argument: {argument}"
+        
+        # Tokenize
+        inputs = self.tokenizer(
+            text,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512,
+            padding=True
+        ).to(self.device)
+        
+        # Predict
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
+            predicted_class = torch.argmax(probabilities, dim=-1).item()
+        
+        # Extract probabilities
+        prob_con = probabilities[0][0].item()
+        prob_pro = probabilities[0][1].item()
+        
+        # Determine stance
+        stance = "PRO" if predicted_class == 1 else "CON"
+        confidence = probabilities[0][predicted_class].item()
         
-    def get_model_info(self):
-        """Get model information"""
-        if not self.model_loaded:
-            return {"loaded": False}
-            
         return {
-            "model_name": self.model_id,
-            "device": str(self.device),
-            "max_length": self.max_length,
-            "num_labels": 2,
-            "loaded": self.model_loaded
+            "predicted_stance": stance,
+            "confidence": confidence,
+            "probability_con": prob_con,
+            "probability_pro": prob_pro
         }
 
 
 # Initialize singleton instance
-kpa_model_manager = KpaModelManager()
+stance_model_manager = StanceModelManager()
+