Update app.py

app.py

@@ -9,40 +9,50 @@ 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 transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, AwqConfig
 from threading import Thread
 import gradio as gr
 
+# Desactivar kernels Marlin incompatibles con el init en CPU de ZeroGPU
+os.environ["GPTQMODEL_DISABLE_MARLIN"] = "1"
+os.environ["AUTOAWQ_USE_MARLIN"] = "0"
+
 MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct-AWQ"
 
 # Cargar Tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
-# Cargar Modelo en AWQ (18-20 GB VRAM)
+# Forzar backend GEMM estándar de AWQ para evitar fallos de repacking en CPU
+quant_config = AwqConfig(
+    bits=4,
+    version="gemm",
+    fuse_max_seq_len=4096,
+    do_fuse=False
+)
+
+# Cargar modelo en modo device_map="auto" compatible con ZeroGPU
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     dtype=torch.float16,
+    quantization_config=quant_config,
     device_map="auto",
     trust_remote_code=True
 )
 
-# Inicializar FastAPI
-app = FastAPI(title="Qwen2.5-Coder-32B OpenAI Compatible API")
+app = FastAPI(title="Qwen2.5-Coder-32B OpenAI API")
 
-# --- Esquemas Pydantic ---
 class ChatMessage(BaseModel):
     role: str
     content: str
 
 class ChatCompletionRequest(BaseModel):
-    model: Optional[str] = "Qwen/Qwen2.5-Coder-32B-Instruct-AWQ"
+    model: Optional[str] = MODEL_ID
     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(
@@ -75,7 +85,6 @@ def generate_stream_tokens(messages_dict: List[Dict[str, str]], temperature: flo
     for new_token in streamer:
         yield new_token
 
-# --- Endpoints OpenAI (/v1) ---
 @app.get("/v1/models")
 async def list_models():
     return {
@@ -96,7 +105,6 @@ async def chat_completions(req: ChatCompletionRequest):
     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(
@@ -110,13 +118,7 @@ async def chat_completions(req: ChatCompletionRequest):
                     "object": "chat.completion.chunk",
                     "created": created_time,
                     "model": req.model,
-                    "choices": [
-                        {
-                            "index": 0,
-                            "delta": {"content": token},
-                            "finish_reason": None
-                        }
-                    ]
+                    "choices": [{"index": 0, "delta": {"content": token}, "finish_reason": None}]
                 }
                 yield f"data: {json.dumps(chunk)}\n\n"
             
@@ -125,20 +127,13 @@ async def chat_completions(req: ChatCompletionRequest):
                 "object": "chat.completion.chunk",
                 "created": created_time,
                 "model": req.model,
-                "choices": [
-                    {
-                        "index": 0,
-                        "delta": {},
-                        "finish_reason": "stop"
-                    }
-                ]
+                "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,
@@ -156,21 +151,13 @@ async def chat_completions(req: ChatCompletionRequest):
         "choices": [
             {
                 "index": 0,
-                "message": {
-                    "role": "assistant",
-                    "content": full_content
-                },
+                "message": {"role": "assistant", "content": full_content},
                 "finish_reason": "stop"
             }
         ],
-        "usage": {
-            "prompt_tokens": -1,
-            "completion_tokens": -1,
-            "total_tokens": -1
-        }
+        "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:
@@ -188,7 +175,7 @@ def gradio_generate(message, history, system_prompt, temperature, top_p, max_tok
 
 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 con interfaz web y endpoints OpenAI (`/v1/chat/completions`).")
+    gr.Markdown("Endpoint OpenAI: `/v1/chat/completions` | Modelo: `Qwen/Qwen2.5-Coder-32B-Instruct-AWQ`")
     gr.ChatInterface(
         fn=gradio_generate,
         additional_inputs=[
@@ -199,7 +186,6 @@ with gr.Blocks(title="Qwen2.5-Coder-32B API & UI", theme=gr.themes.Soft()) as gr
         ]
     )
 
-# Montar Gradio en la raíz
 app = gr.mount_gradio_app(app, gradio_app, path="/")
 
 if __name__ == "__main__":