Update app.py

app.py

@@ -1,151 +1,57 @@
-import os
-import time
-import json
-import uuid
+import spaces
 import gradio as gr
+import torch
 from fastapi import FastAPI, Request
-from fastapi.responses import StreamingResponse, JSONResponse
-from pydantic import BaseModel
-from typing import List, Optional
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 
-from vllm.engine.async_llm_engine import AsyncLLMEngine
-from vllm.engine.arg_utils import AsyncEngineArgs
-from vllm.sampling_params import SamplingParams
-from transformers import AutoTokenizer
+MODEL_ID = "Qwen/Qwen3.8-27B"  # Modelos >14B suelen exceder la memoria dinámica de ZeroGPU
 
-# --- Configuración del Modelo ---
-MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen3.8-27B")
-
-# Argumentos del motor vLLM
-engine_args = AsyncEngineArgs(
-    model=MODEL_NAME,
-    trust_remote_code=True,
-    gpu_memory_utilization=0.90,
-    max_model_len=4096,
-    tensor_parallel_size=1
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
 )
 
-# Inicializar motor y tokenizador
-engine = AsyncLLMEngine.from_engine_args(engine_args)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-
-app = FastAPI(title="vLLM Gradio + OpenAI Server")
-
-# --- Esquemas Pydantic OpenAI ---
-class ChatMessage(BaseModel):
-    role: str
-    content: str
+app = FastAPI()
 
-class ChatCompletionRequest(BaseModel):
-    model: Optional[str] = MODEL_NAME
-    messages: List[ChatMessage]
-    temperature: Optional[float] = 0.7
-    top_p: Optional[float] = 0.95
-    max_tokens: Optional[int] = 2048
-    stream: Optional[bool] = False
-
-# --- Endpoints OpenAI (/v1) ---
-@app.get("/v1/models")
-async def list_models():
-    return {
-        "object": "list",
-        "data": [{"id": MODEL_NAME, "object": "model", "owned_by": "vllm"}]
-    }
-
-@app.post("/v1/chat/completions")
-async def chat_completions(req: ChatCompletionRequest):
-    raw_messages = [{"role": m.role, "content": m.content} for m in req.messages]
-    
-    # Aplicar plantilla de chat del modelo
-    prompt = tokenizer.apply_chat_template(raw_messages, tokenize=False, add_generation_prompt=True)
+@spaces.GPU(duration=120)
+def generate_response(prompt: str, max_tokens: int = 2048):
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_tokens)
     
-    sampling_params = SamplingParams(
-        temperature=req.temperature if req.temperature > 0 else 0.0,
-        top_p=req.top_p,
-        max_tokens=req.max_tokens or 2048
-    )
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
     
-    request_id = f"chatcmpl-{uuid.uuid4().hex}"
+    output_text = ""
+    for new_text in streamer:
+        output_text += new_text
+    return output_text
 
-    if req.stream:
-        async def stream_generator():
-            prev_text = ""
-            results_generator = engine.generate(prompt, sampling_params, request_id)
-            async for request_output in results_generator:
-                text = request_output.outputs[0].text
-                delta = text[len(prev_text):]
-                prev_text = text
-                
-                chunk_data = {
-                    "id": request_id,
-                    "object": "chat.completion.chunk",
-                    "created": int(time.time()),
-                    "model": MODEL_NAME,
-                    "choices": [{"index": 0, "delta": {"content": delta}, "finish_reason": None}]
-                }
-                yield f"data: {json.dumps(chunk_data)}\n\n"
-            yield "data: [DONE]\n\n"
-
-        return StreamingResponse(stream_generator(), media_type="text/event-stream")
-    else:
-        results_generator = engine.generate(prompt, sampling_params, request_id)
-        final_output = None
-        async for request_output in results_generator:
-            final_output = request_output
-            
-        response_text = final_output.outputs[0].text if final_output else ""
-        return {
-            "id": request_id,
-            "object": "chat.completion",
-            "created": int(time.time()),
-            "model": MODEL_NAME,
-            "choices": [{
-                "index": 0,
-                "message": {"role": "assistant", "content": response_text},
-                "finish_reason": "stop"
-            }]
-        }
-
-# --- Interfaz de Gradio ---
-async def gradio_chat(message, history, system_prompt, temperature, 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})
-        messages.append({"role": "assistant", "content": bot_msg})
-        
-    messages.append({"role": "user", "content": message})
-    
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request):
+    data = await request.json()
+    messages = data.get("messages", [])
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    sampling_params = SamplingParams(
-        temperature=temperature if temperature > 0 else 0.0,
-        max_tokens=max_tokens
-    )
     
-    request_id = f"gradio-{uuid.uuid4().hex}"
-    results_generator = engine.generate(prompt, sampling_params, request_id)
+    response_text = generate_response(prompt)
     
-    async for request_output in results_generator:
-        yield request_output.outputs[0].text
-
-with gr.Blocks(title="vLLM + OpenAI Compatible Server") as demo:
-    gr.Markdown(f"# 🚀 vLLM LLM Engine\n**Modelo cargado:** `{MODEL_NAME}`\n\n**Base URL para Cursor / OpenCode:** `https://huggingface.co/spaces/acsaco.hf.space/v1`")
-    
-    with gr.Accordion("Configuración Avanzada", open=False):
-        system_prompt = gr.Textbox(label="System Prompt", value="Eres un asistente experto en programación.")
-        temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.05, label="Temperature")
-        max_tokens = gr.Slider(minimum=128, maximum=4096, value=2048, step=128, label="Max Tokens")
-        
-    gr.ChatInterface(
-        fn=gradio_chat,
-        additional_inputs=[system_prompt, temperature, max_tokens]
-    )
+    return {
+        "id": "chatcmpl-zerogpu",
+        "object": "chat.completion",
+        "choices": [{
+            "index": 0,
+            "message": {
+                "role": "assistant",
+                "content": response_text
+            },
+            "finish_reason": "stop"
+        }]
+    }
 
-# Montar interfaz web Gradio dentro del servidor FastAPI
-app = gr.mount_gradio_app(app, demo, path="/")
+with gr.Blocks() as demo:
+    gr.Markdown("# Qwen ZeroGPU Endpoint")
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+app = gr.mount_gradio_app(app, demo, path="/")