Update Dockerfile

Dockerfile

@@ -1,17 +1,21 @@
-FROM python:3.10-slim
-
-ENV PYTHONDONTWRITEBYTECODE=1 \
-    PYTHONUNBUFFERED=1 \
-    PIP_NO_CACHE_DIR=1 \
-    PORT=7860 \
-    SPACE_E4B_URL=https://cloudunity-gemma-e4b.hf.space \
-    SPACE_MOE_URL=https://cloudunity-gemma-moe26b.hf.space \
-    SPACE_DENSE_URL=https://cloudunity-gemma-dense31b.hf.space
+FROM python:3.11-slim
 
 WORKDIR /app
-RUN apt-get update && apt-get install -y --no-install-recommends git curl build-essential && rm -rf /var/lib/apt/lists/*
-COPY requirements.txt /app/requirements.txt
-RUN pip install --upgrade pip && pip install -r /app/requirements.txt
-COPY app.py /app/app.py
+
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY requirements_simple.txt requirements.txt
+
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY app_simple.py app.py
+
 EXPOSE 7860
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+
+# Set model (change this to use different Gemma 4 variant)
+ENV MODEL_NAME=google/gemma-4-E4B-it
+
+CMD ["python", "app.py"]

app.py

@@ -1,20 +1,22 @@
 import os
 import logging
-import asyncio
-import time
-import uuid
-from typing import Optional, List
+from typing import Optional
 from datetime import datetime
+import time
 
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
-import httpx
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
+# Logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = FastAPI(title="Gemma Orchestrator API - OpenAI Compatible")
+app = FastAPI(title="Gemma 4 Inference API")
+
+# CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -23,186 +25,166 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-SPACES = {
-    "gemma-4-e4b": {
-        "url": os.getenv("SPACE_E4B_URL", "https://example-gemma-e4b.hf.space"),
-        "model": "gemma-4-E4B-it",
-        "size_gb": 5,
-        "context": 128000,
-        "best_for": ["fast", "multimodal", "edge"],
-        "latency_ms": 0,
-        "health": "unknown",
-    },
-    "gemma-4-26b-moe": {
-        "url": os.getenv("SPACE_MOE_URL", "https://example-gemma-moe26b.hf.space"),
-        "model": "gemma-4-26B-A4B-it",
-        "size_gb": 18,
-        "context": 262144,
-        "best_for": ["balanced", "reasoning", "efficient"],
-        "latency_ms": 0,
-        "health": "unknown",
-    },
-    "gemma-4-31b": {
-        "url": os.getenv("SPACE_DENSE_URL", "https://example-gemma-dense31b.hf.space"),
-        "model": "gemma-4-31B-it",
-        "size_gb": 20,
-        "context": 262144,
-        "best_for": ["complex", "coding", "long-context", "reasoning"],
-        "latency_ms": 0,
-        "health": "unknown",
-    },
-}
-
-class ChatCompletionMessage(BaseModel):
+# Config
+MODEL_NAME = os.getenv("MODEL_NAME", "google/gemma-4-E4B-it")
+
+# Global model/tokenizer
+model = None
+tokenizer = None
+
+# Pydantic models
+class Message(BaseModel):
     role: str
     content: str
 
-class ChatCompletionRequest(BaseModel):
-    model: Optional[str] = None
-    messages: List[ChatCompletionMessage]
+class ChatRequest(BaseModel):
+    messages: list[Message]
     temperature: float = Field(default=0.7, ge=0.0, le=2.0)
-    max_tokens: Optional[int] = Field(default=512, ge=1)
+    max_tokens: int = Field(default=512, ge=1, le=2048)
     top_p: float = Field(default=0.9, ge=0.0, le=1.0)
-    top_k: int = Field(default=50, ge=1)
-    stream: bool = False
-    n: int = 1
 
-class ChatCompletionChoice(BaseModel):
+class ChatChoice(BaseModel):
     index: int
-    message: ChatCompletionMessage
+    message: Message
     finish_reason: str
 
-class ChatCompletionUsage(BaseModel):
-    prompt_tokens: int
+class ChatUsage(BaseModel):
     completion_tokens: int
     total_tokens: int
 
-class ChatCompletionResponse(BaseModel):
-    id: str
-    object: str = "chat.completion"
-    created: int
+class ChatResponse(BaseModel):
     model: str
-    choices: List[ChatCompletionChoice]
-    usage: ChatCompletionUsage
-
-class Model(BaseModel):
-    id: str
-    object: str = "model"
+    object: str = "chat.completion"
     created: int
-    owned_by: str = "gemma"
-
-class ModelListResponse(BaseModel):
-    object: str = "list"
-    data: List[Model]
-
-class HealthResponse(BaseModel):
-    status: str
-    spaces: dict
-
-def select_model(model_name: Optional[str], messages: List[ChatCompletionMessage]) -> str:
-    if model_name and model_name in SPACES:
-        return model_name
-    prompt_text = " ".join(msg.content for msg in messages).lower()
-    has_code = any(k in prompt_text for k in ["code", "python", "javascript", "function", "algorithm", "def ", "class "])
-    has_complex = any(k in prompt_text for k in ["reasoning", "explain", "analyze", "compare", "detailed", "comprehensive"])
-    is_short = sum(len(msg.content) for msg in messages) < 500
-    best = "gemma-4-31b" if (has_code or has_complex) else ("gemma-4-e4b" if is_short else "gemma-4-26b-moe")
-    if SPACES[best]["health"] != "healthy":
-        for space_name, info in SPACES.items():
-            if info["health"] == "healthy":
-                return space_name
-    return best
-
-async def check_space_health(space_name: str) -> bool:
-    space = SPACES[space_name]
-    try:
-        async with httpx.AsyncClient(timeout=10) as client:
-            response = await client.get(f"{space['url']}/health")
-        if response.status_code == 200:
-            SPACES[space_name]["health"] = "healthy"
-            return True
-    except Exception:
-        pass
-    SPACES[space_name]["health"] = "unhealthy"
-    return False
-
-async def forward_request(space_name: str, messages: List[ChatCompletionMessage], temperature: float, max_tokens: int, top_p: float, top_k: int):
-    space_url = SPACES[space_name]["url"]
-    payload = {
-        "messages": [{"role": m.role, "content": m.content} for m in messages],
-        "temperature": temperature,
-        "max_tokens": max_tokens,
-        "top_p": top_p,
-        "top_k": top_k,
-    }
-    try:
-        async with httpx.AsyncClient(timeout=120) as client:
-            start = time.time()
-            response = await client.post(f"{space_url}/infer", json=payload)
-            latency = (time.time() - start) * 1000
-            SPACES[space_name]["latency_ms"] = latency
-        response.raise_for_status()
-        return response.json(), latency
-    except httpx.TimeoutException:
-        raise HTTPException(status_code=504, detail=f"Space {space_name} timeout")
-    except Exception as e:
-        raise HTTPException(status_code=502, detail=f"Failed to reach space {space_name}: {e}")
+    choices: list[ChatChoice]
+    usage: ChatUsage
+
+def load_model():
+    """Load model and tokenizer on startup."""
+    global model, tokenizer
+    
+    logger.info(f"Loading {MODEL_NAME}...")
+    
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    
+    # Load with 4-bit quantization to fit in 16GB
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        load_in_4bit=True,
+        low_cpu_mem_usage=True,
+    )
+    
+    logger.info(f"✓ {MODEL_NAME} loaded successfully")
 
 @app.on_event("startup")
-async def startup_event():
-    await asyncio.gather(*[check_space_health(name) for name in SPACES.keys()])
-
-@app.get("/health", response_model=HealthResponse)
-async def health_check():
-    await asyncio.gather(*[check_space_health(name) for name in SPACES.keys()])
-    return HealthResponse(status="healthy", spaces={name: info["health"] for name, info in SPACES.items()})
-
-@app.get("/v1/models", response_model=ModelListResponse)
-@app.get("/models", response_model=ModelListResponse)
-async def list_models():
-    now = int(datetime.utcnow().timestamp())
-    return ModelListResponse(data=[Model(id=model_id, created=now) for model_id in SPACES.keys()])
-
-@app.post("/v1/chat/completions")
-@app.post("/chat/completions")
-async def chat_completions(request: ChatCompletionRequest):
-    selected_model = select_model(request.model, request.messages)
-    response_data, _latency = await forward_request(
-        selected_model,
-        request.messages,
-        request.temperature,
-        request.max_tokens or 512,
-        request.top_p,
-        request.top_k,
-    )
-    completion_id = str(uuid.uuid4())
-    tokens = int(response_data.get("tokens_used", 0))
-    return ChatCompletionResponse(
-        id=f"chatcmpl-{completion_id}",
-        created=int(datetime.utcnow().timestamp()),
-        model=selected_model,
-        choices=[ChatCompletionChoice(index=0, message=ChatCompletionMessage(role="assistant", content=response_data["response"]), finish_reason="stop")],
-        usage=ChatCompletionUsage(prompt_tokens=0, completion_tokens=tokens, total_tokens=tokens),
-    )
+async def startup():
+    load_model()
+
+@app.get("/health")
+def health():
+    return {"status": "ok", "model": MODEL_NAME}
 
-@app.get("/info")
-async def orchestrator_info():
+@app.get("/v1/models")
+def list_models():
     return {
-        "orchestrator": "Gemma Orchestrator - OpenAI Compatible",
-        "version": "1.0",
-        "spaces": {name: {k: v for k, v in info.items() if k != "url"} for name, info in SPACES.items()},
+        "object": "list",
+        "data": [
+            {
+                "id": "gemma-4",
+                "object": "model",
+                "owned_by": "google",
+                "created": int(time.time()),
+            }
+        ]
     }
 
+@app.post("/v1/chat/completions", response_model=ChatResponse)
+def chat_completions(request: ChatRequest):
+    """OpenAI-compatible chat completions endpoint."""
+    
+    try:
+        # Build prompt from messages
+        prompt = ""
+        for msg in request.messages:
+            if msg.role == "system":
+                prompt += f"<|system|>\n{msg.content}<|end_of_turn|>\n"
+            elif msg.role == "user":
+                prompt += f"<|user|>\n{msg.content}<|end_of_turn|>\n"
+            elif msg.role == "assistant":
+                prompt += f"<|assistant|>\n{msg.content}<|end_of_turn|>\n"
+        
+        prompt += "<|assistant|>\n"
+        
+        # Tokenize
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        input_length = inputs.input_ids.shape[1]
+        
+        # Generate
+        start_time = time.time()
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=request.max_tokens,
+            temperature=request.temperature,
+            top_p=request.top_p,
+            do_sample=request.temperature > 0,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+        
+        # Decode
+        full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract just the response
+        if "<|assistant|>" in full_text:
+            response_text = full_text.split("<|assistant|>")[-1].strip()
+        else:
+            response_text = full_text
+        
+        tokens_generated = outputs.shape[1] - input_length
+        
+        return ChatResponse(
+            model="gemma-4",
+            created=int(time.time()),
+            choices=[
+                ChatChoice(
+                    index=0,
+                    message=Message(role="assistant", content=response_text),
+                    finish_reason="stop",
+                )
+            ],
+            usage=ChatUsage(
+                completion_tokens=tokens_generated,
+                total_tokens=tokens_generated,
+            ),
+        )
+    
+    except Exception as e:
+        logger.error(f"Error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/chat/completions", response_model=ChatResponse)
+def chat_completions_no_v1(request: ChatRequest):
+    """Alias without /v1/ prefix."""
+    return chat_completions(request)
+
 @app.get("/")
-async def root():
+def root():
     return {
-        "name": "Gemma Orchestrator API",
-        "version": "1.0",
-        "openai_compatible": True,
-        "endpoints": {
-            "/v1/models": "List models",
-            "/v1/chat/completions": "Chat completions",
-            "/health": "Health check",
-            "/info": "Orchestrator info"
+        "name": "Gemma 4 API",
+        "model": MODEL_NAME,
+        "docs": "Use /v1/chat/completions for OpenAI compatibility",
+        "example": {
+            "url": "/v1/chat/completions",
+            "method": "POST",
+            "body": {
+                "messages": [{"role": "user", "content": "Hello"}],
+                "temperature": 0.7,
+                "max_tokens": 512,
+            }
         }
     }
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -1,14 +1,9 @@
-fastapi
-uvicorn
-python-multipart
-pydantic
-requests
-httpx
-torch
-transformers
-huggingface-hub
-numpy
-python-dotenv
-sentencepiece
-accelerate
-openai
+fastapi==0.115.6
+uvicorn==0.32.1
+pydantic==2.10.5
+pydantic-settings==2.7.1
+torch==2.6.0 --index-url https://download.pytorch.org/whl/cpu
+transformers==4.48.1
+huggingface-hub==0.26.5
+bitsandbytes==0.44.1
+accelerate==1.2.1