Update app.py

app.py

@@ -7,12 +7,9 @@ import time
 import uuid
 import json
 from typing import Optional, List, Union, Dict, Any
-import asyncio
-
-# --- LA SEULE LIGNE À MODIFIER ---
-MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct" # Ou "deepseek-ai/deepseek-coder-1.3b-instruct", etc.
-# ------------------------------------
 
+# --- Configuration ---
+MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
 DEVICE = "cpu"
 
 # --- Chargement du modèle ---
@@ -23,18 +20,12 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map=DEVICE
 )
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-
-# --- Standardisation : Gestion automatique du padding ---
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-    print("Le pad_token a été défini sur eos_token.")
-
 print("Modèle et tokenizer chargés avec succès sur le CPU.")
 
 # --- Création de l'application API ---
 app = FastAPI()
 
-# --- Modèles de données standards pour la compatibilité OpenAI ---
+# --- Modèles de données pour accepter la structure complexe de l'extension ---
 class ContentPart(BaseModel):
     type: str
     text: str
@@ -51,18 +42,6 @@ class ChatCompletionRequest(BaseModel):
     class Config:
         extra = Extra.ignore
 
-class ChatCompletionResponseChoice(BaseModel):
-    index: int = 0
-    message: ChatMessage
-    finish_reason: str = "stop"
-
-class ChatCompletionResponse(BaseModel):
-    id: str
-    object: str = "chat.completion"
-    created: int
-    model: str
-    choices: List[ChatCompletionResponseChoice]
-
 class ModelData(BaseModel):
     id: str
     object: str = "model"
@@ -71,7 +50,7 @@ class ModelData(BaseModel):
 class ModelList(BaseModel):
     object: str = "list"
     data: List[ModelData]
-    
+
 # --- Définition des API ---
 
 @app.get("/models", response_model=ModelList)
@@ -83,36 +62,74 @@ async def list_models():
 async def create_chat_completion(request: ChatCompletionRequest):
     """Endpoint principal qui gère la génération de texte en streaming."""
     
-    # --- Standardisation : On ne fait aucune supposition sur le format du message ---
-    # On passe la liste complète des messages au tokenizer.
-    messages_for_model = [msg.dict() for msg in request.messages]
-
-    # La fonction apply_chat_template gère automatiquement le "dialecte" de chaque modèle.
-    text_prompt = tokenizer.apply_chat_template(messages_for_model, tokenize=False, add_generation_prompt=True)
-    
-    inputs = tokenizer(text_prompt, return_tensors="pt", padding=True).to(DEVICE)
+    # On extrait le prompt de l'utilisateur de la structure complexe
+    user_prompt = ""
+    last_message = request.messages[-1]
+    if isinstance(last_message.content, list):
+        for part in last_message.content:
+            if part.type == 'text':
+                user_prompt += part.text + "\n"
+    elif isinstance(last_message.content, str):
+        user_prompt = last_message.content
+
+    if not user_prompt:
+        return {"error": "Prompt non trouvé."}
+
+    # Préparation pour le modèle DeepSeek
+    messages_for_model = [{'role': 'user', 'content': user_prompt}]
+    inputs = tokenizer.apply_chat_template(messages_for_model, add_generation_prompt=True, return_tensors="pt").to(DEVICE)
     
-    outputs = model.generate(**inputs, max_new_tokens=250, do_sample=True, temperature=0.2, top_k=50, top_p=0.95, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
-    
-    response_text = tokenizer.decode(outputs[0, inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+    # Génération de la réponse complète
+    outputs = model.generate(inputs, max_new_tokens=250, do_sample=True, temperature=0.2, top_k=50, top_p=0.95, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
+    response_text = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
 
+    # Fonction génératrice pour le streaming
     async def stream_generator():
         response_id = f"chatcmpl-{uuid.uuid4()}"
         
+        # On envoie la réponse caractère par caractère, au format attendu
         for char in response_text:
-            chunk = { "id": response_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": MODEL_ID, "choices": [{"index": 0, "delta": {"content": char}, "finish_reason": None }] }
+            chunk = {
+                "id": response_id,
+                "object": "chat.completion.chunk",
+                "created": int(time.time()),
+                "model": MODEL_ID,
+                "choices": [{
+                    "index": 0,
+                    "delta": {"content": char},
+                    "finish_reason": None
+                }]
+            }
             yield f"data: {json.dumps(chunk)}\n\n"
-            await asyncio.sleep(0.01)
+            await asyncio.sleep(0.01) # Petite pause pour simuler un flux
         
-        final_chunk = { "id": response_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": MODEL_ID, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop" }] }
+        # On envoie le chunk final de fin
+        final_chunk = {
+            "id": response_id,
+            "object": "chat.completion.chunk",
+            "created": int(time.time()),
+            "model": MODEL_ID,
+            "choices": [{
+                "index": 0,
+                "delta": {},
+                "finish_reason": "stop"
+            }]
+        }
         yield f"data: {json.dumps(final_chunk)}\n\n"
+        
+        # On envoie le signal [DONE]
         yield "data: [DONE]\n\n"
 
+    # Si l'extension demande un stream, on renvoie le générateur
     if request.stream:
         return StreamingResponse(stream_generator(), media_type="text/event-stream")
     else:
+        # Code de secours si le stream n'est pas demandé (peu probable)
         return {"choices": [{"message": {"role": "assistant", "content": response_text}}]}
 
 @app.get("/")
 def root():
-    return {"status": "API compatible OpenAI en ligne (avec streaming)", "model_id": MODEL_ID}
+    return {"status": "API compatible OpenAI en ligne (avec streaming)", "model_id": MODEL_ID}
+
+# On a besoin de asyncio pour la pause dans le stream
+import asyncio