Update app.py

app.py

@@ -1,19 +1,19 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from pydantic import BaseModel
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import time
+import uuid
 
 # --- Configuration ---
-# ON UTILISE LE MODÈLE DEEPSEEK CODER 1.3B INSTRUCT
 MODEL_ID = "deepseek-ai/deepseek-coder-1.3b-instruct"
 DEVICE = "cpu"
 
 # --- Chargement du modèle et du tokenizer ---
-# Cette étape peut prendre quelques minutes au premier démarrage du Space
 print(f"Début du chargement du modèle : {MODEL_ID}")
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.bfloat16, # bfloat16 est efficace sur CPU
+    torch_dtype=torch.bfloat16,
     device_map=DEVICE
 )
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -22,33 +22,64 @@ print("Modèle et tokenizer chargés avec succès sur le CPU.")
 # --- Création de l'application API ---
 app = FastAPI()
 
-# Modèle de données pour la requête d'entrée
-class PromptRequest(BaseModel):
-    prompt: str
-    max_new_tokens: int = 250
+# --- Modèles de données pour la compatibilité OpenAI ---
+class ChatMessage(BaseModel):
+    role: str
+    content: str
 
-# --- Définition de l'API ---
-@app.post("/generate")
-async def generate_code(request: PromptRequest):
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: list[ChatMessage]
+    max_tokens: int = 250
+
+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]
+
+# --- Définition de l'API compatible OpenAI ---
+@app.post("/v1/chat/completions")
+async def create_chat_completion(request: ChatCompletionRequest):
     """
-    Endpoint pour générer du code à partir d'un prompt.
+    Endpoint compatible avec l'API OpenAI Chat Completions.
     """
-    # Le modèle DeepSeek Instruct utilise un format de chat spécifique
-    messages = [
-        {'role': 'user', 'content': request.prompt}
-    ]
+    # Extraire le dernier message utilisateur pour le prompt
+    user_prompt = ""
+    if request.messages and request.messages[-1].role == "user":
+        user_prompt = request.messages[-1].content
     
-    # Préparation des inputs pour le modèle
-    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(DEVICE)
+    if not user_prompt:
+        return {"error": "No user prompt found"}
+
+    # Préparation des inputs 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)
     
-    # Génération du code. Soyez patient, cela peut prendre plusieurs secondes sur CPU.
-    outputs = model.generate(inputs, max_new_tokens=request.max_new_tokens, do_sample=True, temperature=0.2, top_k=50, top_p=0.95, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
+    # Génération
+    outputs = model.generate(inputs, max_new_tokens=request.max_tokens, do_sample=True, temperature=0.2, top_k=50, top_p=0.95, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
     
-    # Décodage de la réponse
-    response = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
+    # Décodage
+    response_text = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
+
+    # Formatage de la réponse au format OpenAI
+    response_message = ChatMessage(role="assistant", content=response_text)
+    choice = ChatCompletionResponseChoice(message=response_message)
+    completion_response = ChatCompletionResponse(
+        id=f"chatcmpl-{uuid.uuid4()}",
+        created=int(time.time()),
+        model=request.model,
+        choices=[choice]
+    )
 
-    return {"generated_text": response}
+    return completion_response
 
 @app.get("/")
 def root():
-    return {"status": "API en ligne", "model_id": MODEL_ID, "message": "Ce modèle tourne sur un CPU gratuit."}
+    return {"status": "API compatible OpenAI en ligne", "model_id": MODEL_ID}