Add OpenAI-compatible API server for distil-home-assistant-functiongemma

Dockerfile

@@ -0,0 +1,16 @@
+FROM python:3.10-slim
+
+RUN useradd -m -u 1000 user
+WORKDIR /app
+
+COPY --chown=user requirements.txt .
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user app.py .
+
+USER user
+ENV HOME=/home/user
+ENV HF_HOME=/home/user/.cache/huggingface
+
+EXPOSE 7860
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,110 @@
+"""
+distil-home-assistant-functiongemma — OpenAI/OpenRouter-compatible inference server
+
+Endpoints:
+    GET  /                       -> health check / API info
+    GET  /v1/models              -> list available models
+    POST /v1/chat/completions    -> generate text (OpenAI format)
+"""
+
+import os
+import time
+import uuid
+
+import torch
+from fastapi import FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# ─── Config ───────────────────────────────────────────────────────────────────
+
+MODEL_ID = "LisaMegaWatts/distil-home-assistant-functiongemma"
+MODEL_NAME = "functiongemma-270m"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
+
+# ─── Load model at startup ────────────────────────────────────────────────────
+
+print(f"Loading {MODEL_ID} on {DEVICE} ...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID, torch_dtype=DTYPE, trust_remote_code=True
+).to(DEVICE)
+model.eval()
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+MODEL_CREATED_AT = int(time.time())
+print(f"Model ready on {DEVICE}")
+
+# ─── FastAPI app ──────────────────────────────────────────────────────────────
+
+app = FastAPI(title=f"{MODEL_NAME} API", version="1.0.0")
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
+
+
+@app.get("/")
+def root():
+    return {
+        "name": MODEL_NAME,
+        "version": "1.0.0",
+        "description": "FunctionGemma 270m distilled for home assistant function calling",
+        "model": MODEL_ID,
+        "endpoints": ["/v1/models", "/v1/chat/completions"],
+        "compatible_with": ["OpenAI API", "OpenRouter"],
+    }
+
+
+@app.get("/v1/models")
+def list_models():
+    return {
+        "object": "list",
+        "data": [{"id": MODEL_NAME, "object": "model", "created": MODEL_CREATED_AT, "owned_by": "LisaMegaWatts"}],
+    }
+
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request):
+    try:
+        body = await request.json()
+    except Exception:
+        return JSONResponse(status_code=400, content={"error": {"message": "Invalid JSON", "type": "invalid_request_error", "code": "invalid_json"}})
+
+    messages = body.get("messages", [])
+    temperature = max(0.01, min(float(body.get("temperature", 0.7)), 2.0))
+    max_tokens = max(1, min(int(body.get("max_tokens", 512)), 2048))
+    top_p = float(body.get("top_p", 0.9))
+    n_completions = int(body.get("n", 1))
+
+    if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template:
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    else:
+        parts = [f"{m.get('role','user')}: {m.get('content','')}" for m in messages]
+        parts.append("assistant:")
+        prompt = "\n".join(parts)
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+    prompt_tokens = inputs["input_ids"].shape[1]
+
+    choices = []
+    total_completion_tokens = 0
+    for i in range(n_completions):
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs, max_new_tokens=max_tokens, temperature=temperature,
+                top_p=top_p, do_sample=temperature > 0.01, pad_token_id=tokenizer.pad_token_id,
+            )
+        new_tokens = outputs[0][prompt_tokens:]
+        text = tokenizer.decode(new_tokens, skip_special_tokens=True)
+        completion_tokens = len(new_tokens)
+        total_completion_tokens += completion_tokens
+        choices.append({"index": i, "message": {"role": "assistant", "content": text}, "finish_reason": "length" if completion_tokens >= max_tokens else "stop"})
+
+    return {
+        "id": f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion", "created": int(time.time()),
+        "model": MODEL_NAME, "choices": choices,
+        "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": total_completion_tokens, "total_tokens": prompt_tokens + total_completion_tokens},
+        "system_fingerprint": f"{MODEL_NAME}-v1",
+    }

requirements.txt

@@ -0,0 +1,9 @@
+fastapi
+uvicorn[standard]
+torch
+transformers
+accelerate
+safetensors
+sentencepiece
+protobuf
+tokenizers