Update app.py

app.py

@@ -2,25 +2,32 @@ import os
 import sys
 import time
 import re
+import json
 import multiprocessing
+from typing import List, Dict, Any, Optional, Union, AsyncGenerator
+
+# FastAPI & Pydantic imports
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
 
 # ---------------------------------------------------------------------------
-# ⚡ 1. ENVIRONMENT TUNING
+# ⚡ 1. CPU & THREADING OPTIMIZATIONS FOR MoE
 # ---------------------------------------------------------------------------
-# Auto-detect physical CPU cores to avoid hyperthreading contention
+# MoE models are heavily bottlenecked by CPU RAM bandwidth during expert routing.
+# Locking threads to physical cores avoids hyperthreading thrashing.
 try:
     import psutil
     PHYSICAL_CORES = psutil.cpu_count(logical=False) or max(1, multiprocessing.cpu_count() // 2)
 except Exception:
     PHYSICAL_CORES = max(1, multiprocessing.cpu_count() // 2)
 
-# Lock thread pools strictly to physical CPU cores
 os.environ["OMP_NUM_THREADS"] = str(PHYSICAL_CORES)
 os.environ["MKL_NUM_THREADS"] = str(PHYSICAL_CORES)
 os.environ["OPENBLAS_NUM_THREADS"] = str(PHYSICAL_CORES)
 os.environ["OMP_PROC_BIND"] = "TRUE"
 
-# Raise OS scheduler priority (if permitted by the environment)
 try:
     p = psutil.Process()
     if sys.platform == "win32":
@@ -30,54 +37,142 @@ try:
 except Exception:
     pass
 
-# Import heavy libraries AFTER environment variables are set
 import gradio as gr
 import llama_cpp
 from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
 
 # ---------------------------------------------------------------------------
-# ⚡ 2. MODEL DOWNLOAD & LLAMA-CPP INITIALIZATION
+# ⚡ 2. MODEL DOWNLOAD & LLAMA-CPP MoE INITIALIZATION
 # ---------------------------------------------------------------------------
 NUMA_STRATEGY = getattr(llama_cpp, "GGML_NUMA_STRATEGY_DISTRIBUTE", False)
 
-print("Downloading model...")
+print("Downloading MoE model GGUF...")
 model_path = hf_hub_download(
     repo_id="unsloth/Qwen3.6-35B-A3B-GGUF",
     filename="Qwen3.6-35B-A3B-UD-IQ3_XXS.gguf",
 )
 print("Model downloaded successfully!")
 
-# Load GGUF model optimized for CPU execution
+# Load GGUF model optimized for CPU MoE execution
 llm = Llama(
     model_path=model_path,
     n_gpu_layers=0,             # 0 for CPU execution
     n_ctx=4096,
-    chat_format="chatml",       # ChatML format prevents KV cache issues
     
     # Threading optimizations
     n_threads=PHYSICAL_CORES,       
     n_threads_batch=PHYSICAL_CORES, 
     numa=NUMA_STRATEGY,
     
-    # Memory & Cache Management
-    use_mlock=False,            # False is safer for Hugging Face Spaces to avoid OOM
-    use_mmap=True,
-    flash_attn=False,           # Disabled on CPU
+    # MoE Cache & RAM Bandwidth Optimizations:
+    # 1. Quantized KV Cache (Q4_0) saves RAM and frees CPU cache space for routing
+    type_k=getattr(llama_cpp, "GGML_TYPE_Q4_0", 2),
+    type_v=getattr(llama_cpp, "GGML_TYPE_Q4_0", 2),
     
-    # Batch Memory Alignment
+    # 2. Smaller ubatch prevents L2/L3 cache eviction during active expert evaluation
     n_batch=256,    
-    n_ubatch=64,   
+    n_ubatch=32,   
+    
+    use_mlock=False,
+    use_mmap=True,
+    flash_attn=False,
 )
 
+MODEL_NAME = "Qwen3.6-35B-A3B"
+
+# ---------------------------------------------------------------------------
+# ⚡ 3. FASTAPI INITIALIZATION & OPENAI PYDANTIC SCHEMAS
+# ---------------------------------------------------------------------------
+app = FastAPI(title="OpenAI-Compatible CPU MoE Server", version="1.0.0")
+
+# Enable CORS for external web frontend requests
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+class ChatMessage(BaseModel):
+    role: str
+    content: Union[str, List[Dict[str, Any]]]
+
+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] = 512
+    stream: Optional[bool] = False
+
+# ---------------------------------------------------------------------------
+# ⚡ 4. OPENAI-COMPATIBLE REST ENDPOINTS
+# ---------------------------------------------------------------------------
+@app.get("/v1/models")
+async def list_models():
+    """OpenAI standard /v1/models endpoint."""
+    return {
+        "object": "list",
+        "data": [
+            {
+                "id": MODEL_NAME,
+                "object": "model",
+                "created": int(time.time()),
+                "owned_by": "local",
+            }
+        ],
+    }
+
+async def openai_stream_generator(messages_dict: list, max_tokens: int, temperature: float, top_p: float):
+    """Generates standard SSE streaming tokens for OpenAI client libraries."""
+    for chunk in llm.create_chat_completion(
+        messages=messages_dict,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        stream=True,
+    ):
+        yield f"data: {json.dumps(chunk)}\n\n"
+    yield "data: [DONE]\n\n"
+
+@app.post("/v1/chat/completions")
+async def chat_completions(req: ChatCompletionRequest):
+    """OpenAI standard /v1/chat/completions endpoint."""
+    try:
+        # Convert pydantic messages to standard list of dicts
+        messages_dict = [{"role": m.role, "content": m.content} for m in req.messages]
+        
+        if req.stream:
+            return StreamingResponse(
+                openai_stream_generator(
+                    messages_dict=messages_dict,
+                    max_tokens=req.max_tokens,
+                    temperature=req.temperature,
+                    top_p=req.top_p,
+                ),
+                media_type="text/event-stream",
+            )
+        else:
+            response = llm.create_chat_completion(
+                messages=messages_dict,
+                max_tokens=req.max_tokens,
+                temperature=req.temperature,
+                top_p=req.top_p,
+                stream=False,
+            )
+            return JSONResponse(content=response)
+            
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
 # ---------------------------------------------------------------------------
-# ⚡ 3. ROBUST HISTORY CLEANER
+# ⚡ 5. GRADIO UI & HISTORY CLEANER
 # ---------------------------------------------------------------------------
 def clean_assistant_reply(text):
     if not text:
         return ""
-    
-    # Safely handle Gradio 5 / Multimodal list & dict formats
     if isinstance(text, list):
         extracted = []
         for part in text:
@@ -89,22 +184,14 @@ def clean_assistant_reply(text):
     elif not isinstance(text, str):
         text = str(text)
 
-    # 1. Strip reasoning blocks (<think>...</think>)
     cleaned = re.sub(r'<think>.*?</think>\n*', '', text, flags=re.DOTALL)
     cleaned = cleaned.replace("</think>", "")
-    
-    # 2. Strip previous timing stats footers
     cleaned = re.sub(r'\n\n---\n\*⏱️.*?\*$', '', cleaned)
-    
     return cleaned.strip()
 
-# ---------------------------------------------------------------------------
-# ⚡ 4. GRADIO RESPOND FUNCTION 
-# ---------------------------------------------------------------------------
 def respond(message, history, system_message, max_tokens, temperature, top_p):
     messages = [{"role": "system", "content": system_message}]
     
-    # Parse chat history safely
     for item in history:
         if isinstance(item, dict):
             role = item.get("role", "user")
@@ -138,7 +225,6 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
             token_count += 1
             yield response
             
-    # Calculate generation speed
     end_time = time.time()
     elapsed_time = end_time - start_time
     
@@ -148,20 +234,24 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         response += stats_footer
         yield response
 
-# ---------------------------------------------------------------------------
-# ⚡ 5. GRADIO INTERFACE LAUNCH
-# ---------------------------------------------------------------------------
 chatbot = gr.ChatInterface(
     fn=respond,
     additional_inputs=[
-        gr.Textbox(value="You are a highly capable coding assistant.", label="System message"),
+        gr.Textbox(value="You are a helpful coding assistant.", label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
     ],
-    title="Qwopus 3.5 Coder (Ultra-Optimized CPU Edition)",
+    title="Qwen 3.6 MoE (FastAPI + Gradio Engine)",
 )
 
+# Mount Gradio UI inside FastAPI at the root path "/"
+app = gr.mount_gradio_app(app, chatbot, path="/")
+
+# ---------------------------------------------------------------------------
+# ⚡ 6. SERVER LAUNCH
+# ---------------------------------------------------------------------------
 if __name__ == "__main__":
-    # Specify the 0.0.0.0 host and port 7860 to ensure HF spaces can access the health check
-    chatbot.launch(server_name="0.0.0.0", server_port=7860)
+    import uvicorn
+    # Port 7860 is required by Hugging Face Spaces for health checks
+    uvicorn.run(app, host="0.0.0.0", port=7860)