Update app.py

app.py

@@ -1,33 +1,58 @@
-# Fichier: app.py (VERSION CORRIGÉE FINALE)
+# Fichier: app.py (VERSION CORRIGÉE FINALE - OPTIMISÉE POUR LA MÉMOIRE DU SPACE)
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig 
 import torch
+import os
 
 # --- Configuration du Modèle ---
-# Nouveau Modèle : Phi-3 Mini (3.8B) - Optimisé pour les systèmes légers/CPU
 model_id = "microsoft/Phi-3-mini-4k-instruct"
-
-# Détecter le périphérique : forcer le CPU
 device = torch.device("cpu")
 
+# --- Stratégie de chargement pour économiser la mémoire (Quantisation) ---
+# Si le Space a un GPU/CUDA, la quantisation sera utilisée, réduisant la RAM par 8.
+# Si le Space est CPU seulement, cette tentative échouera, et nous utiliserons le fallback float32.
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+)
+
 # Charger le Tokenizer et le Modèle
 try:
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     
-    # CORRECTION CRITIQUE 1: Stabilité du chargement sur CPU
-    # 1. Ajout de torch_dtype=torch.float32 pour assurer la compatibilité CPU.
-    # 2. Suppression de device_map=device (le .to(device) final est suffisant).
+    # TENTATIVE 1 : Chargement avec Quantisation 4-bit (Méthode recommandée)
+    print("Tentative de chargement avec quantisation 4-bit...")
+    # Le chargement en 4-bit nécessite device_map="auto"
     model = AutoModelForCausalLM.from_pretrained(
         model_id, 
-        torch_dtype=torch.float32,
-        trust_remote_code=True # Nécessaire pour Phi-3
-    ).to(device)
-    model.eval()
-    print(f"Modèle {model_id} chargé sur CPU.")
-except Exception as e:
-    print(f"Erreur lors du chargement du modèle : {e}")
-    raise e
+        quantization_config=quantization_config,
+        device_map="auto",
+        trust_remote_code=True
+    )
+    print(f"Modèle {model_id} chargé et quantifié.")
+
+except Exception as e_quant:
+    # Si la quantisation échoue (souvent sans GPU), on revient à la version CPU
+    print(f"Échec de la quantisation : {e_quant}. Tentative de chargement float32 CPU (Attention: peut causer OOM).")
+    
+    # TENTATIVE 2 : Fallback sur le chargement float32 CPU (Votre code initial, mais avec fix du bug)
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id, 
+            torch_dtype=torch.float32, 
+            trust_remote_code=True
+        ).to(device)
+        print(f"Modèle {model_id} chargé sur CPU (Float32).")
+    except Exception as e_cpu:
+        print(f"Échec critique du chargement CPU : {e_cpu}")
+        # Si même float32 échoue, vous avez BESOIN de plus de RAM pour votre Space.
+        raise e_cpu
+
+model.eval()
 
 app = FastAPI(
     title="NLP Space - Phi-3 Mini API (CPU)",
@@ -48,12 +73,12 @@ def generate_text_from_model(system_prompt: str, user_prompt: str, max_tokens: i
         {"role": "user", "content": user_prompt}
     ]
     
-    # Appliquer le template de chat du tokenizer
     text_to_generate = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
     
-    # Exemple de format: "<|user|>\nInstruction\n<|end|>\n<|assistant|>"
-    
-    inputs = tokenizer(text_to_generate, return_tensors="pt").to(device)
+    # Trouver le device réel du modèle pour y placer les inputs (nécessaire après le chargement device_map)
+    # Assurez-vous que le modèle est correctement placé, en le forçant sur CPU si nécessaire.
+    real_device = model.device if model.device.type != 'meta' else torch.device("cpu")
+    inputs = tokenizer(text_to_generate, return_tensors="pt").to(real_device)
     
     with torch.no_grad():
         output = model.generate(
@@ -62,13 +87,12 @@ def generate_text_from_model(system_prompt: str, user_prompt: str, max_tokens: i
             do_sample=True, 
             temperature=temperature,
             pad_token_id=tokenizer.eos_token_id,
-            # CORRECTION CRITIQUE 2: Désactiver le cache pour contourner le bug 'DynamicCache'
-            use_cache=False 
+            use_cache=False # CORRECTION CRITIQUE 2: Fixe le bug DynamicCache
         )
     
     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
     
-    # Nettoyage : retirer l'instruction initiale pour ne garder que la réponse
+    # Nettoyage
     response_start_tag = "<|assistant|>"
     if response_start_tag in generated_text:
         return generated_text.split(response_start_tag, 1)[1].strip()
@@ -76,7 +100,7 @@ def generate_text_from_model(system_prompt: str, user_prompt: str, max_tokens: i
     return generated_text.strip()
 
 
-# --- Endpoints (Identiques) ---
+# --- Endpoints (Inchangés) ---
 
 @app.post("/generate")
 async def generate(request: PromptRequest):
@@ -91,12 +115,11 @@ async def generate(request: PromptRequest):
         )
         return {"result": result}
     except Exception as e:
-        # Retourne l'erreur Python pour le diagnostic (comme vous l'avez fait)
         return {"error": str(e)}
 
 @app.post("/summarize")
 async def summarize(request: PromptRequest):
-    """Résumé d'un long texte."""
+    # ... (code inchangé) ...
     system_prompt = "Tu es un expert en résumé concis et précis. Ton objectif est de résumer le texte fourni de manière à en conserver l'idée principale."
     user_prompt = f"Résume le texte suivant de manière concise et factuelle:\n\n---\n\n{request.prompt}"
     try:
@@ -112,7 +135,7 @@ async def summarize(request: PromptRequest):
 
 @app.post("/classify")
 async def classify(request: PromptRequest):
-    """Classification du sentiment, du thème ou de la catégorie d'un texte."""
+    # ... (code inchangé) ...
     system_prompt = "Tu es un expert en classification. Réponds uniquement avec l'étiquette de classification demandée sans phrases supplémentaires."
     user_prompt = request.prompt
     try: