Upload step03_chatbot.py with huggingface_hub

step03_chatbot.py

@@ -9,6 +9,18 @@ import json
 import numpy as np
 import gradio as gr
 
+# Import spaces pour ZeroGPU compatibility
+try:
+    import spaces
+except ImportError:
+    # Fallback pour environnements non-ZeroGPU
+    class spaces:
+        @staticmethod
+        def GPU(duration=60):
+            def decorator(func):
+                return func
+            return decorator
+
 # Désactiver le warning tokenizers sur ZeroGPU
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 from gradio import ChatMessage
@@ -156,7 +168,7 @@ class Qwen3Reranker:
     Reranker utilisant Qwen3-Reranker-4B pour améliorer la pertinence des résultats de recherche
     """
     
-    def __init__(self, model_name: str = "Qwen/Qwen3-Reranker-4B", use_flash_attention: bool = True):
+    def __init__(self, model_name: str = "Qwen/Qwen3-Reranker-4B", use_flash_attention: bool = False):
         """
         Initialise le reranker Qwen3
         
@@ -165,7 +177,7 @@ class Qwen3Reranker:
             use_flash_attention: Utiliser Flash Attention 2 si disponible (auto-désactivé sur Mac)
         """
         self.model_name = model_name
-        self.use_flash_attention = use_flash_attention
+        self.use_flash_attention = False  # Désactivé pour éviter les problèmes
         
         # Détection de l'environnement
         self.is_mps = torch.backends.mps.is_available()
@@ -258,15 +270,15 @@ class Qwen3Reranker:
         if self.is_mps:
             self.device = torch.device("mps")
             self.model = self.model.to(self.device)
-        elif self.is_cuda and not os.getenv("SPACE_ID"):
-            # Utiliser CUDA seulement si pas sur ZeroGPU Spaces
+        elif self.is_cuda:
+            # Utiliser CUDA si disponible
             if hasattr(self.model, 'device'):
                 self.device = next(self.model.parameters()).device
             else:
                 self.device = torch.device("cuda")
                 self.model = self.model.to(self.device)
         else:
-            # Forcer CPU sur ZeroGPU pour éviter l'erreur CUDA init
+            # Fallback CPU
             self.device = torch.device("cpu")
             self.model = self.model.to(self.device)
     
@@ -407,7 +419,7 @@ class GenericRAGChatbot:
                  generation_model: str = "Qwen/Qwen3-4B-Instruct-2507",
                  initial_k: int = 20,
                  final_k: int = 3,
-                 use_flash_attention: bool = True,
+                 use_flash_attention: bool = False,
                  use_reranker: bool = True):
         """
         Initialise le système RAG générique
@@ -422,7 +434,7 @@ class GenericRAGChatbot:
         self.generation_model_name = generation_model
         self.initial_k = initial_k
         self.final_k = final_k
-        self.use_flash_attention = use_flash_attention
+        self.use_flash_attention = False  # Désactivé pour éviter les problèmes
         self.use_reranker = use_reranker
         
         # Détection de l'environnement (local + ZeroGPU)
@@ -571,18 +583,7 @@ class GenericRAGChatbot:
         try:
             from sentence_transformers import SentenceTransformer
             
-            if os.getenv("SPACE_ID"):
-                print("  - Configuration ZeroGPU optimisée")
-                # Sur ZeroGPU, utiliser float16 et device auto pour les performances
-                self.embedding_model = SentenceTransformer(
-                    self.config.embedding_model,
-                    model_kwargs={
-                        "torch_dtype": torch.float16,
-                        "device_map": "auto"
-                    },
-                    tokenizer_kwargs={"padding_side": "left"}
-                )
-            elif self.use_flash_attention and self.is_cuda:
+            if self.use_flash_attention and self.is_cuda:
                 print("  - Configuration avec Flash Attention 2 activée (CUDA)")
                 try:
                     self.embedding_model = SentenceTransformer(
@@ -722,7 +723,7 @@ class GenericRAGChatbot:
         except:
             return 0.0
 
-    @spaces.GPU(duration=120)  # ZeroGPU: GPU nécessaire pour embedding
+    @spaces.GPU(duration=120)
     def search_documents(self, query: str, final_k: int = None, use_reranking: bool = None) -> List[Dict]:
         """
         Recherche avancée avec reranking en deux étapes
@@ -819,7 +820,7 @@ class GenericRAGChatbot:
         
         return final_results
 
-    @spaces.GPU(duration=120)  # ZeroGPU: GPU seulement pour la génération
+    @spaces.GPU(duration=180)
     def generate_response_stream(self, query: str, context: str, history: List = None):
         """
         Génère une réponse streamée basée sur le contexte et l'historique
@@ -913,6 +914,7 @@ Instructions importantes:
         except Exception as e:
             yield f"❌ Erreur lors de la génération: {str(e)}"
 
+    @spaces.GPU(duration=180)
     def generate_response(self, query: str, context: str, history: List = None) -> str:
         """
         Génère une réponse basée sur le contexte et l'historique
@@ -997,6 +999,7 @@ Réponds à cette question en te basant sur le contexte fourni."""
             print(f"❌ Erreur lors de la génération: {e}")
             return f"❌ Erreur lors de la génération de la réponse: {str(e)}"
 
+    @spaces.GPU(duration=300)  # Durée plus longue car combine search + generation
     def stream_response_with_tools(self, query: str, history, top_k: int = None, use_reranking: bool = None):
         """
         Génère une réponse streamée avec affichage visuel des tools et reranking Qwen3
@@ -1156,7 +1159,7 @@ def _create_rag_system():
     if is_zerogpu:
         default_config = {
             'generation_model': "Qwen/Qwen3-4B-Instruct-2507",    # Modèle qui fonctionne sur ZeroGPU
-            'use_flash_attention': True,   # ZeroGPU supporte Flash Attention
+            'use_flash_attention': False,  # Désactivé pour stabilité
             'use_reranker': True,          # GPU puissant, reranking activé
             'initial_k': 20,               # Même config que local
             'final_k': 5                   # Plus de documents finaux
@@ -1170,7 +1173,7 @@ def _create_rag_system():
         }
     else:
         default_config = {
-            'use_flash_attention': is_cuda,  # Flash Attention seulement sur CUDA
+            'use_flash_attention': False,    # Désactivé pour stabilité
             'use_reranker': True,            # Reranking par défaut
             'initial_k': 20,                 # Candidats pour la première étape
             'final_k': 3                     # Documents finaux par défaut
@@ -1205,6 +1208,7 @@ def _ensure_chatmessages(history):
     return result
 
 
+@spaces.GPU(duration=300)  # Fonction principale de chat
 def chat_with_generic_rag(message, history, top_k, use_reranking):
     """
     Interface entre Gradio et le système RAG générique avec contrôles avancés.