Initial deploy

Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.10-slim
+
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1 \
+    PORT=7860 \
+    MODEL_NAME=gemma-4-E4B-it \
+    MODEL_SIZE=5 \
+    MODEL_CONTEXT=128000
+
+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
+EXPOSE 7860
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,172 @@
+import os
+import logging
+import time
+from typing import Optional
+from datetime import datetime
+from functools import lru_cache
+
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+MODEL_NAME = os.getenv("MODEL_NAME", "gemma-4-E4B-it")
+MODEL_SIZE = int(os.getenv("MODEL_SIZE", "5"))
+MODEL_CONTEXT = int(os.getenv("MODEL_CONTEXT", "128000"))
+HF_REPO_ID = os.getenv("HF_REPO_ID", f"google/{MODEL_NAME}")
+QUANTIZATION = os.getenv("QUANTIZATION", "Q4_K_M")
+
+DEVICE = "cpu"
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+    DEVICE = "mps"
+
+app = FastAPI(title=f"Gemma Inference API - {MODEL_NAME}")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class InferenceRequest(BaseModel):
+    messages: list[ChatMessage]
+    model: str = MODEL_NAME
+    temperature: float = 0.7
+    max_tokens: int = 512
+    top_p: float = 0.9
+    top_k: int = 50
+    thinking: bool = False
+    system_prompt: Optional[str] = None
+
+class InferenceResponse(BaseModel):
+    model: str
+    response: str
+    tokens_used: int
+    latency_ms: float
+    thinking: Optional[str] = None
+    timestamp: str
+
+class HealthResponse(BaseModel):
+    status: str
+    model: str
+    device: str
+    model_size_gb: int
+    context_window: int
+
+@lru_cache(maxsize=1)
+def get_tokenizer():
+    logger.info(f"Loading tokenizer for {HF_REPO_ID}...")
+    return AutoTokenizer.from_pretrained(HF_REPO_ID)
+
+@lru_cache(maxsize=1)
+def get_model():
+    logger.info(f"Loading model {HF_REPO_ID} on {DEVICE}...")
+    kwargs = dict(low_cpu_mem_usage=True)
+    if DEVICE == "cuda":
+        kwargs["torch_dtype"] = torch.float16
+        model = AutoModelForCausalLM.from_pretrained(HF_REPO_ID, device_map="auto", **kwargs)
+    elif DEVICE == "cpu" and MODEL_SIZE > 10:
+        kwargs["torch_dtype"] = torch.float32
+        try:
+            model = AutoModelForCausalLM.from_pretrained(HF_REPO_ID, device_map="auto", load_in_4bit=True, **kwargs)
+        except Exception:
+            model = AutoModelForCausalLM.from_pretrained(HF_REPO_ID, **kwargs)
+    else:
+        kwargs["torch_dtype"] = torch.float32
+        model = AutoModelForCausalLM.from_pretrained(HF_REPO_ID, **kwargs)
+    return model
+
+def build_chat_prompt(messages: list[ChatMessage], system_prompt: Optional[str] = None) -> str:
+    parts = []
+    if system_prompt:
+        parts.append(f"<|system|>\n{system_prompt}<|end_of_turn|>\n")
+    for msg in messages:
+        parts.append(f"<|{msg.role}|>\n{msg.content}<|end_of_turn|>\n")
+    parts.append("<|assistant|>\n")
+    return "".join(parts)
+
+@app.get("/health", response_model=HealthResponse)
+async def health_check():
+    return HealthResponse(status="healthy", model=MODEL_NAME, device=DEVICE, model_size_gb=MODEL_SIZE, context_window=MODEL_CONTEXT)
+
+@app.get("/info")
+async def model_info():
+    return {"model": MODEL_NAME, "repo_id": HF_REPO_ID, "device": DEVICE, "model_size_gb": MODEL_SIZE, "context_window": MODEL_CONTEXT, "quantization": QUANTIZATION}
+
+@app.post("/infer", response_model=InferenceResponse)
+async def infer(request: InferenceRequest):
+    start_time = time.time()
+    try:
+        tokenizer = get_tokenizer()
+        model = get_model()
+        prompt = build_chat_prompt(request.messages, system_prompt=request.system_prompt)
+        if request.thinking:
+            prompt = "<|think|>\n" + prompt
+        inputs = tokenizer(prompt, return_tensors="pt")
+        model_device = getattr(model, 'device', None)
+        if model_device is not None:
+            inputs = {k: v.to(model_device) for k, v in inputs.items()}
+        input_length = inputs['input_ids'].shape[1]
+        max_new = max(1, min(request.max_tokens, MODEL_CONTEXT - input_length))
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new,
+            temperature=request.temperature,
+            top_p=request.top_p,
+            top_k=request.top_k,
+            do_sample=request.temperature > 0,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response_text = generated_text.split("<|assistant|>")[-1].strip() if "<|assistant|>" in generated_text else generated_text
+        latency_ms = (time.time() - start_time) * 1000
+        tokens_generated = outputs.shape[1] - input_length
+        return InferenceResponse(model=MODEL_NAME, response=response_text, tokens_used=int(tokens_generated), latency_ms=latency_ms, timestamp=datetime.utcnow().isoformat())
+    except Exception as e:
+        logger.exception("Inference error")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/chat")
+async def chat(request: InferenceRequest):
+    return await infer(request)
+
+@app.post("/complete")
+async def complete(request: InferenceRequest):
+    if not request.messages:
+        raise HTTPException(status_code=400, detail="No messages provided")
+    simple_request = InferenceRequest(
+        messages=[ChatMessage(role="user", content=request.messages[-1].content)],
+        temperature=request.temperature,
+        max_tokens=request.max_tokens,
+        top_p=request.top_p,
+        top_k=request.top_k,
+        thinking=request.thinking,
+        system_prompt=request.system_prompt,
+    )
+    return await infer(simple_request)
+
+@app.get("/")
+async def root():
+    return {
+        "name": "Gemma Inference API",
+        "model": MODEL_NAME,
+        "version": "1.0",
+        "endpoints": {
+            "/health": "Health check",
+            "/info": "Model information",
+            "/infer": "Run inference (POST)",
+            "/chat": "Chat interface (POST)",
+            "/complete": "Text completion (POST)"
+        }
+    }

requirements.txt

@@ -0,0 +1,15 @@
+fastapi==0.104.1
+uvicorn==0.24.0
+python-multipart==0.0.6
+pydantic==2.5.0
+requests==2.31.0
+httpx==0.25.2
+torch==2.1.2
+transformers==4.36.2
+huggingface-hub==0.20.1
+numpy==1.26.3
+python-dotenv==1.0.0
+sentencepiece==0.1.99
+accelerate==0.25.0
+bitsandbytes==0.42.0
+openai==1.3.7