Create app.py

app.py

@@ -0,0 +1,435 @@
+import os
+import threading
+import time
+from collections import deque
+
+import torch
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+
+# ============================================================
+# CONFIG
+# ============================================================
+
+MODEL_PATH = os.getenv(
+    "MODEL_PATH",
+    "./gemma3-270m-merged"
+)
+
+MAX_INPUT_TOKENS = int(
+    os.getenv("MAX_INPUT_TOKENS", "1024")
+)
+
+MAX_NEW_TOKENS = int(
+    os.getenv("MAX_NEW_TOKENS", "256")
+)
+
+MAX_HISTORY_MESSAGES = int(
+    os.getenv("MAX_HISTORY_MESSAGES", "8")
+)
+
+# Number of simultaneous generations.
+# Keep this LOW on CPU.
+MAX_CONCURRENT_GENERATIONS = int(
+    os.getenv("MAX_CONCURRENT_GENERATIONS", "1")
+)
+
+# Maximum waiting requests.
+MAX_QUEUE_SIZE = int(
+    os.getenv("MAX_QUEUE_SIZE", "20")
+)
+
+# CPU threads.
+CPU_THREADS = int(
+    os.getenv(
+        "CPU_THREADS",
+        str(max(1, (os.cpu_count() or 4) - 1))
+    )
+)
+
+torch.set_num_threads(CPU_THREADS)
+
+DEVICE = "cpu"
+
+
+# ============================================================
+# APP
+# ============================================================
+
+app = FastAPI(
+    title="Gemma 3 270M API",
+    description="CPU inference API for Gemma 3 270M",
+    version="1.0.0",
+)
+
+
+# ============================================================
+# REQUEST / RESPONSE MODELS
+# ============================================================
+
+class Message(BaseModel):
+    role: str
+    content: str
+
+
+class ChatRequest(BaseModel):
+    message: str = Field(
+        ...,
+        min_length=1,
+        max_length=12000
+    )
+
+    history: list[Message] = Field(
+        default_factory=list
+    )
+
+    max_new_tokens: int = Field(
+        default=128,
+        ge=1,
+        le=256
+    )
+
+    temperature: float = Field(
+        default=0.7,
+        ge=0.0,
+        le=2.0
+    )
+
+    top_p: float = Field(
+        default=0.9,
+        gt=0.0,
+        le=1.0
+    )
+
+
+class ChatResponse(BaseModel):
+    response: str
+    model: str
+    input_tokens: int
+    output_tokens: int
+    generation_time: float
+
+
+# ============================================================
+# GLOBAL STATE
+# ============================================================
+
+print("=" * 70)
+print("             GEMMA 3 270M CPU API SERVER")
+print("=" * 70)
+
+print(f"Model: {MODEL_PATH}")
+print(f"Device: {DEVICE}")
+print(f"CPU threads: {CPU_THREADS}")
+
+print()
+print("Loading tokenizer...")
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_PATH,
+    local_files_only=True,
+)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+print("Tokenizer loaded.")
+
+print()
+print("Loading model...")
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_PATH,
+    local_files_only=True,
+    dtype=torch.float32,
+    low_cpu_mem_usage=True,
+)
+
+model.to(DEVICE)
+model.eval()
+
+print("Model loaded successfully.")
+print()
+
+# Semaphore prevents multiple CPU generations from hammering
+# the machine simultaneously.
+generation_semaphore = threading.BoundedSemaphore(
+    MAX_CONCURRENT_GENERATIONS
+)
+
+# Simple queue counter.
+queue_lock = threading.Lock()
+waiting_requests = 0
+
+
+# ============================================================
+# HEALTH
+# ============================================================
+
+@app.get("/")
+def root():
+    return {
+        "name": "Gemma 3 270M API",
+        "status": "online",
+        "model": "Gemma 3 270M",
+        "device": "cpu",
+        "api": "/v1/chat",
+    }
+
+
+@app.get("/health")
+def health():
+    return {
+        "status": "healthy",
+        "model_loaded": True,
+        "device": DEVICE,
+        "cpu_threads": CPU_THREADS,
+        "waiting_requests": waiting_requests,
+    }
+
+
+# ============================================================
+# CHAT
+# ============================================================
+
+@app.post(
+    "/v1/chat",
+    response_model=ChatResponse
+)
+def chat(request: ChatRequest):
+
+    global waiting_requests
+
+    # --------------------------------------------------------
+    # Queue protection
+    # --------------------------------------------------------
+
+    with queue_lock:
+
+        if waiting_requests >= MAX_QUEUE_SIZE:
+            raise HTTPException(
+                status_code=429,
+                detail=(
+                    "Server is busy. "
+                    "Please try again later."
+                )
+            )
+
+        waiting_requests += 1
+
+    acquired = False
+
+    try:
+
+        # ----------------------------------------------------
+        # Wait for generation slot
+        # ----------------------------------------------------
+
+        generation_semaphore.acquire()
+        acquired = True
+
+        # ----------------------------------------------------
+        # Prepare conversation
+        # ----------------------------------------------------
+
+        messages = []
+
+        history = request.history[
+            -MAX_HISTORY_MESSAGES:
+        ]
+
+        for item in history:
+
+            if item.role not in (
+                "user",
+                "assistant"
+            ):
+                continue
+
+            content = item.content.strip()
+
+            if not content:
+                continue
+
+            messages.append(
+                {
+                    "role": item.role,
+                    "content": content,
+                }
+            )
+
+        messages.append(
+            {
+                "role": "user",
+                "content": request.message.strip(),
+            }
+        )
+
+        # ----------------------------------------------------
+        # Gemma chat template
+        # ----------------------------------------------------
+
+        try:
+
+            prompt = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+
+        except Exception:
+
+            # Fallback if tokenizer template is unavailable
+            prompt = request.message.strip()
+
+        # ----------------------------------------------------
+        # Tokenize
+        # ----------------------------------------------------
+
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=MAX_INPUT_TOKENS,
+        )
+
+        inputs = {
+            key: value.to(DEVICE)
+            for key, value in inputs.items()
+        }
+
+        input_tokens = inputs[
+            "input_ids"
+        ].shape[1]
+
+        # ----------------------------------------------------
+        # Generation
+        # ----------------------------------------------------
+
+        max_tokens = min(
+            request.max_new_tokens,
+            MAX_NEW_TOKENS
+        )
+
+        start_time = time.perf_counter()
+
+        with torch.inference_mode():
+
+            if request.temperature <= 0:
+
+                outputs = model.generate(
+                    **inputs,
+
+                    max_new_tokens=max_tokens,
+
+                    do_sample=False,
+
+                    pad_token_id=tokenizer.pad_token_id,
+                    eos_token_id=tokenizer.eos_token_id,
+
+                    use_cache=True,
+                )
+
+            else:
+
+                outputs = model.generate(
+                    **inputs,
+
+                    max_new_tokens=max_tokens,
+
+                    do_sample=True,
+
+                    temperature=request.temperature,
+                    top_p=request.top_p,
+
+                    repetition_penalty=1.10,
+
+                    pad_token_id=tokenizer.pad_token_id,
+                    eos_token_id=tokenizer.eos_token_id,
+
+                    use_cache=True,
+                )
+
+        generation_time = (
+            time.perf_counter() - start_time
+        )
+
+        # ----------------------------------------------------
+        # Decode ONLY generated tokens
+        # ----------------------------------------------------
+
+        generated_tokens = outputs[
+            0,
+            input_tokens:
+        ]
+
+        response = tokenizer.decode(
+            generated_tokens,
+            skip_special_tokens=True,
+        ).strip()
+
+        if not response:
+            response = "I couldn't generate a response."
+
+        output_tokens = generated_tokens.shape[0]
+
+        return ChatResponse(
+            response=response,
+            model="gemma-3-270m",
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            generation_time=round(
+                generation_time,
+                3
+            ),
+        )
+
+    finally:
+
+        if acquired:
+            generation_semaphore.release()
+
+        with queue_lock:
+            waiting_requests = max(
+                0,
+                waiting_requests - 1
+            )
+
+
+# ============================================================
+# STARTUP MESSAGE
+# ============================================================
+
+if __name__ == "__main__":
+
+    import uvicorn
+
+    print("=" * 70)
+    print("SERVER READY")
+    print("=" * 70)
+
+    print()
+    print("API:")
+    print("POST /v1/chat")
+
+    print()
+    print("Health:")
+    print("GET /health")
+
+    print()
+    print("Swagger:")
+    print("GET /docs")
+
+    print()
+    print("Starting server...")
+
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=int(
+            os.getenv("PORT", "7860")
+        ),
+        workers=1,
+    )