app.py

from fastapi import FastAPI
from pydantic import BaseModel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# --- 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
    device_map=DEVICE
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
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

# --- Définition de l'API ---
@app.post("/generate")
async def generate_code(request: PromptRequest):
    """
    Endpoint pour générer du code à partir d'un prompt.
    """
    # Le modèle DeepSeek Instruct utilise un format de chat spécifique
    messages = [
        {'role': 'user', 'content': request.prompt}
    ]
    
    # Préparation des inputs pour le modèle
    inputs = tokenizer.apply_chat_template(messages, 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)
    
    # Décodage de la réponse
    response = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)

    return {"generated_text": response}

@app.get("/")
def root():
    return {"status": "API en ligne", "model_id": MODEL_ID, "message": "Ce modèle tourne sur un CPU gratuit."}