Update app.py

app.py

@@ -1,16 +1,24 @@
 import os
+import time
+import json
+import uuid
 import torch
 import spaces
-import gradio as gr
+import uvicorn
+from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+from typing import List, Optional, Dict
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
+import gradio as gr
 
 MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct-AWQ"
 
-# Cargar tokenizer
+# Cargar Tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
-# Cargar modelo cuantizado en AWQ (ocupa ~18-20GB VRAM)
+# Cargar Modelo en AWQ (18-20 GB VRAM)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.float16,
@@ -18,21 +26,27 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True
 )
 
-@spaces.GPU(duration=120)
-def generate(message, history, system_prompt, temperature, top_p, max_tokens):
-    messages = []
-    if system_prompt:
-        messages.append({"role": "system", "content": system_prompt})
+# Inicializar FastAPI
+app = FastAPI(title="Qwen2.5-Coder-32B OpenAI Compatible API")
 
-    for user_msg, bot_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        if bot_msg:
-            messages.append({"role": "assistant", "content": bot_msg})
+# --- Esquemas Pydantic ---
+class ChatMessage(BaseModel):
+    role: str
+    content: str
 
-    messages.append({"role": "user", "content": message})
+class ChatCompletionRequest(BaseModel):
+    model: Optional[str] = "Qwen/Qwen2.5-Coder-32B-Instruct-AWQ"
+    messages: List[ChatMessage]
+    temperature: Optional[float] = 0.2
+    top_p: Optional[float] = 0.9
+    max_tokens: Optional[int] = 2048
+    stream: Optional[bool] = False
 
+# Función central de generación protegida por ZeroGPU
+@spaces.GPU(duration=120)
+def generate_stream_tokens(messages_dict: List[Dict[str, str]], temperature: float, top_p: float, max_tokens: int):
     text = tokenizer.apply_chat_template(
-        messages,
+        messages_dict,
         tokenize=False,
         add_generation_prompt=True
     )
@@ -48,58 +62,145 @@ def generate(message, history, system_prompt, temperature, top_p, max_tokens):
     generate_kwargs = dict(
         model_inputs,
         streamer=streamer,
-        max_new_tokens=int(max_tokens),
+        max_new_tokens=max_tokens,
         do_sample=temperature > 0.0,
         temperature=max(temperature, 1e-2) if temperature > 0.0 else None,
         top_p=top_p if temperature > 0.0 else None,
         repetition_penalty=1.05
     )
 
-    # Iniciar generación en un hilo secundario para streaming fluido
     thread = Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
 
-    partial_text = ""
     for new_token in streamer:
-        partial_text += new_token
+        yield new_token
+
+# --- Endpoints OpenAI (/v1) ---
+@app.get("/v1/models")
+async def list_models():
+    return {
+        "object": "list",
+        "data": [
+            {
+                "id": MODEL_ID,
+                "object": "model",
+                "created": int(time.time()),
+                "owned_by": "huggingface"
+            }
+        ]
+    }
+
+@app.post("/v1/chat/completions")
+async def chat_completions(req: ChatCompletionRequest):
+    req_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
+    created_time = int(time.time())
+    messages_dict = [{"role": m.role, "content": m.content} for m in req.messages]
+
+    # Manejo de streaming (SSE)
+    if req.stream:
+        async def event_generator():
+            for token in generate_stream_tokens(
+                messages_dict,
+                temperature=req.temperature or 0.2,
+                top_p=req.top_p or 0.9,
+                max_tokens=req.max_tokens or 2048
+            ):
+                chunk = {
+                    "id": req_id,
+                    "object": "chat.completion.chunk",
+                    "created": created_time,
+                    "model": req.model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"content": token},
+                            "finish_reason": None
+                        }
+                    ]
+                }
+                yield f"data: {json.dumps(chunk)}\n\n"
+            
+            final_chunk = {
+                "id": req_id,
+                "object": "chat.completion.chunk",
+                "created": created_time,
+                "model": req.model,
+                "choices": [
+                    {
+                        "index": 0,
+                        "delta": {},
+                        "finish_reason": "stop"
+                    }
+                ]
+            }
+            yield f"data: {json.dumps(final_chunk)}\n\n"
+            yield "data: [DONE]\n\n"
+
+        return StreamingResponse(event_generator(), media_type="text/event-stream")
+
+    # Respuesta estándar (sin streaming)
+    full_content = ""
+    for token in generate_stream_tokens(
+        messages_dict,
+        temperature=req.temperature or 0.2,
+        top_p=req.top_p or 0.9,
+        max_tokens=req.max_tokens or 2048
+    ):
+        full_content += token
+
+    return {
+        "id": req_id,
+        "object": "chat.completion",
+        "created": created_time,
+        "model": req.model,
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": full_content
+                },
+                "finish_reason": "stop"
+            }
+        ],
+        "usage": {
+            "prompt_tokens": -1,
+            "completion_tokens": -1,
+            "total_tokens": -1
+        }
+    }
+
+# --- Interfaz Gradio ---
+def gradio_generate(message, history, system_prompt, temperature, top_p, max_tokens):
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+    for user_msg, bot_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        if bot_msg:
+            messages.append({"role": "assistant", "content": bot_msg})
+    messages.append({"role": "user", "content": message})
+
+    partial_text = ""
+    for token in generate_stream_tokens(messages, temperature, top_p, int(max_tokens)):
+        partial_text += token
         yield partial_text
 
-# Interfaz de Gradio
-with gr.Blocks(title="Qwen2.5-Coder-32B-Instruct AWQ (ZeroGPU)", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="Qwen2.5-Coder-32B API & UI", theme=gr.themes.Soft()) as gradio_app:
     gr.Markdown("# 🚀 Qwen2.5-Coder-32B-Instruct (AWQ en ZeroGPU)")
-    gr.Markdown("Servicio de inferencia optimizado para código y compatible con clientes vía API / Gradio.")
-
-    chat_interface = gr.ChatInterface(
-        fn=generate,
+    gr.Markdown("Servicio con interfaz web y endpoints OpenAI (`/v1/chat/completions`).")
+    gr.ChatInterface(
+        fn=gradio_generate,
         additional_inputs=[
-            gr.Textbox(
-                value="Eres un asistente de programación experto, preciso y conciso.",
-                label="System Prompt",
-                lines=2
-            ),
-            gr.Slider(
-                minimum=0.0,
-                maximum=1.0,
-                value=0.2,
-                step=0.05,
-                label="Temperature"
-            ),
-            gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=0.9,
-                step=0.05,
-                label="Top-P"
-            ),
-            gr.Slider(
-                minimum=256,
-                maximum=4096,
-                value=2048,
-                step=256,
-                label="Max New Tokens"
-            )
+            gr.Textbox("Eres un asistente de programación experto.", label="System Prompt"),
+            gr.Slider(0.0, 1.0, 0.2, step=0.05, label="Temperature"),
+            gr.Slider(0.1, 1.0, 0.9, step=0.05, label="Top-P"),
+            gr.Slider(256, 4096, 2048, step=256, label="Max New Tokens")
         ]
     )
 
+# Montar Gradio en la raíz
+app = gr.mount_gradio_app(app, gradio_app, path="/")
+
 if __name__ == "__main__":
-    demo.queue().launch()
+    uvicorn.run(app, host="0.0.0.0", port=7860)