Enhance KpaModelManager to load fine-tuned weights from Hugging Face and update requirements to include huggingface_hub

requirements.txt

@@ -6,4 +6,5 @@ torch>=2.0.0
 transformers>=4.35.0
 accelerate>=0.24.0
 protobuf>=3.20.0
+huggingface_hub>=0.19.0
 

services/label_model_manage.py

@@ -3,6 +3,7 @@
 import os
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from huggingface_hub import hf_hub_download
 import logging
 
 logger = logging.getLogger(__name__)
@@ -20,7 +21,7 @@ class KpaModelManager:
         self.model_id = None
 
     def load_model(self, model_id: str, api_key: str = None):
-        """Load model and tokenizer from Hugging Face"""
+        """Load model with weights from Hugging Face repository"""
         if self.model_loaded:
             logger.info("KPA model already loaded")
             return
@@ -38,20 +39,35 @@ class KpaModelManager:
             # Prepare token for authentication if API key is provided
             token = api_key if api_key else None
             
-            # Load tokenizer and model from Hugging Face
-            logger.info("Loading tokenizer...")
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                model_id,
-                token=token,
-                trust_remote_code=True
-            )
+            # 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)
             
-            logger.info("Loading model...")
+            # Load base model architecture
+            logger.info(f"Loading base model architecture from {base_model_name}...")
             self.model = AutoModelForSequenceClassification.from_pretrained(
-                model_id,
-                token=token,
-                trust_remote_code=True
+                base_model_name,
+                num_labels=2
+            )
+            
+            # 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(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()