Upload step03_chatbot.py with huggingface_hub

step03_chatbot.py

@@ -334,6 +334,11 @@ class Qwen3Reranker:
             print("  - Reranker non disponible, scores neutres retournés")
             return [0.5] * len(documents)
         
+        # Assurer que le modèle est sur GPU dans ZeroGPU
+        if torch.cuda.is_available() and not next(self.model.parameters()).is_cuda:
+            print("  - Déplacement du reranker vers GPU...")
+            self.model = self.model.cuda()
+        
         if instruction is None:
             instruction = self._get_default_instruction()
         
@@ -708,6 +713,13 @@ class GenericRAGChatbot:
         
         print(f"🔍 Recherche en deux étapes: {initial_k} candidats → reranking → {k} finaux")
         
+        # Assurer que le modèle embedding est sur GPU dans ZeroGPU
+        if torch.cuda.is_available() and hasattr(self.embedding_model, 'model'):
+            device = next(self.embedding_model.model.parameters()).device
+            if not device.type == 'cuda':
+                print("  - Déplacement du modèle embedding vers GPU...")
+                self.embedding_model = self.embedding_model.cuda()
+        
         # Étape 1: Recherche par embedding avec FAISS
         if hasattr(self.embedding_model, 'prompts') and 'query' in self.embedding_model.prompts:
             query_embedding = self.embedding_model.encode([query], prompt_name="query")[0]
@@ -803,6 +815,11 @@ class GenericRAGChatbot:
             yield "❌ Modèle de génération non disponible"
             return
         
+        # Assurer que le modèle de génération est sur GPU dans ZeroGPU
+        if torch.cuda.is_available() and not next(self.generation_model.parameters()).is_cuda:
+            print("  - Déplacement du modèle de génération vers GPU...")
+            self.generation_model = self.generation_model.cuda()
+        
         # Construction du prompt système
         system_prompt = """Tu es un assistant expert qui répond aux questions en te basant uniquement sur les documents fournis dans le contexte.
 
@@ -1110,7 +1127,7 @@ def _create_rag_system():
             'generation_model': "Qwen/Qwen3-4B-Instruct-2507",    # Modèle qui fonctionne sur ZeroGPU
             'use_flash_attention': True,   # ZeroGPU supporte Flash Attention
             'use_reranker': True,          # GPU puissant, reranking activé
-            'initial_k': 30,               # Plus de candidats avec GPU puissant
+            'initial_k': 20,               # Même config que local
             'final_k': 5                   # Plus de documents finaux
         }
     elif is_mac: