Upload app.py

app.py

@@ -1,32 +1,29 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import uvicorn
-import os
-
-import traceback
-
-from fastapi.middleware.cors import CORSMiddleware
+from threading import Thread
+import json
 
 app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_credentials=False, # Plus sûr avec wildcards
+    allow_credentials=False,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
 # --- CONFIGURATION ---
-VERSION = "2.3 (Final Access Fix)"
+VERSION = "3.0 (Streaming)"
 MODEL_ID = "issoufzousko07/BABA-IA-2B"
 
-print("="*50)
 print(f"🚀 Démarrage BABA API v{VERSION}")
 print(f"Chargement de {MODEL_ID}...")
-print("="*50)
 
 # Détection automatique
 if torch.cuda.is_available():
@@ -39,8 +36,7 @@ if torch.cuda.is_available():
     )
 else:
     device = "cpu"
-    print("🐢 Mode CPU activé (Gemma 2B tient dans la RAM)")
-    # Sur CPU, on évite device_map="auto" pour éviter les bugs d'offloading accelerate
+    print("🐢 Mode CPU activé (Streaming activé pour compenser la lenteur)")
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
         torch_dtype=torch.float32
@@ -55,45 +51,41 @@ class ChatRequest(BaseModel):
 
 @app.post("/chat")
 async def chat(request: ChatRequest):
-    try:
-        print(f"📩 Reçu : {request.message}")
-        
-        # 1. MESSAGE + TEMPLATE
-        messages = [{"role": "user", "content": request.message}]
-        
-        # 1. TEMPLATE -> TEXTE (Plus sûr)
-        # On récupère le prompt complet sous forme de string
-        text_prompt = tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True
-        )
-        
-        # 2. TEXTE -> TENSEURS
-        # On tokenise explicitement pour avoir input_ids ET attention_mask
-        inputs = tokenizer(text_prompt, return_tensors="pt").to(model.device)
+    print(f"📩 Reçu (Stream) : {request.message}")
+    
+    # 1. MESSAGE + TEMPLATE
+    messages = [{"role": "user", "content": request.message}]
+    text_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(text_prompt, return_tensors="pt").to(model.device)
 
-        # 3. GÉNÉRATION
-        # On passe **inputs pour envoyer input_ids + attention_mask correctement
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=300,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9,
-        )
+    # 2. CONFIG STREAMER
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    
+    generation_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=300,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+    )
 
-        # 4. DÉCODAGE
-        # inputs.input_ids.shape[-1] donne la longueur du prompt
-        response = tokenizer.decode(outputs[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True)
-        print(f"📤 Réponse : {response.strip()[:50]}...")
-        
-        return {"response": response.strip()}
+    # 3. GÉNÉRATION DANS UN THREAD (Non-bloquant)
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    # 4. GÉNÉRATEUR DE RÉPONSE SSE
+    def stream_response():
+        full_text = ""
+        for token in streamer:
+            full_text += token
+            # Format SSE: data: {"token": "..."}\n\n
+            yield f"data: {json.dumps({'token': token, 'text': full_text})}\n\n"
         
-    except Exception as e:
-        print("❌ ERREUR CRITIQUE :")
-        traceback.print_exc() # Affiche toute l'erreur dans les logs
-        return {"response": f"Erreur technique : {str(e)}"}
+        # Signal de fin
+        yield "data: [DONE]\n\n"
+
+    return StreamingResponse(stream_response(), media_type="text/event-stream")
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)