Update app.py

app.py

@@ -1,101 +1,77 @@
-import os
-import torch
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
 import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 
 # --- Configurações ---
 MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct"
 
-# Variáveis globais iniciadas como None (vazias)
+print(f"⏳ Iniciando carregamento preguiçoso para {MODEL_ID}...")
+
+# Variáveis globais para cache do modelo
 model = None
 tokenizer = None
 
-app = FastAPI(title="APIDOST - Gabriel's Router")
-
-# --- Estruturas de Dados ---
-class Message(BaseModel):
-    role: str
-    content: str
-
-class ChatCompletionRequest(BaseModel):
-    model: str = "default-model"
-    messages: list[Message]
-    max_tokens: int = 1024
-    temperature: float = 0.7
-
-# --- A Mágica do ZeroGPU com Lazy Loading ---
-# duration=120 garante 2 minutos de GPU, tempo suficiente pro load + inferencia
-@spaces.GPU(duration=120) 
-def generate_response(messages, max_tokens, temperature):
+def load_model():
     global model, tokenizer
-    
-    # O PULO DO GATO: Só carrega se ainda não estiver na memória
     if model is None:
-        print(f"⏳ Cold Start: Carregando {MODEL_ID} para a VRAM...")
+        print("🚀 Carregando modelo para a VRAM (Cold Start)...")
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
             torch_dtype=torch.bfloat16,
-            device_map="cuda" # Força o uso da GPU alocada pelo spaces
+            device_map="cuda" # O @spaces.GPU garante que 'cuda' é a H200
         )
-        print("✅ Modelo carregado com sucesso!")
+        print("✅ Modelo carregado!")
+    return model, tokenizer
+
+# --- A Função Mágica do ZeroGPU ---
+@spaces.GPU(duration=120)
+def respond(message, history):
+    # Carrega o modelo apenas quando a GPU é alocada
+    model, tokenizer = load_model()
+
+    # Formata o histórico para o padrão do Qwen
+    messages = []
+    for user_msg, bot_msg in history:
+        if user_msg: messages.append({"role": "user", "content": user_msg})
+        if bot_msg: messages.append({"role": "assistant", "content": bot_msg})
+    
+    messages.append({"role": "user", "content": message})
 
     # Prepara o prompt
-    text_prompt = tokenizer.apply_chat_template(
-        messages, 
-        tokenize=False, 
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
         add_generation_prompt=True
     )
     
-    inputs = tokenizer(text_prompt, return_tensors="pt").to(model.device)
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-    # Gera
-    output = model.generate(
-        **inputs,
-        max_new_tokens=max_tokens,
-        temperature=temperature,
-        do_sample=True,
-        top_p=0.9,
+    # Configuração de geração
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=1024,
+        temperature=0.7,
+        do_sample=True
     )
     
-    response_text = tokenizer.decode(output[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
-    return response_text
-
-# --- Endpoints ---
-@app.get("/")
-def read_root():
-    return {"status": "APIDOST is online", "mode": "Lazy Loading Active"}
-
-@app.post("/v1/chat/completions")
-async def chat_completions(request: ChatCompletionRequest):
-    try:
-        # Converte mensagens
-        msgs = [{"role": m.role, "content": m.content} for m in request.messages]
-        
-        # Chama a função protegida pelo @spaces.GPU
-        response_content = generate_response(msgs, request.max_tokens, request.temperature)
-
-        return {
-            "id": "chatcmpl-apidost",
-            "object": "chat.completion",
-            "created": 1234567890,
-            "model": request.model,
-            "choices": [{
-                "index": 0,
-                "message": {
-                    "role": "assistant",
-                    "content": response_content
-                },
-                "finish_reason": "stop"
-            }]
-        }
+    # Decodifica a resposta
+    generated_ids = [
+        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    ]
+    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    
+    return response
 
-    except Exception as e:
-        print(f"❌ Erro: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+# --- Interface Gradio ---
+demo = gr.ChatInterface(
+    respond,
+    title="APIDOST - Qwen 2.5 Coder (H200 Powered)",
+    description="Rodando no ZeroGPU da Hugging Face. Use via API ou Chat.",
+    examples=["Crie um script Python para snake game.", "Explique a teoria da relatividade."],
+)
 
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    demo.launch()