import os import logging from contextlib import asynccontextmanager from typing import List, Optional import torch from fastapi import FastAPI, HTTPException from fastapi.responses import HTMLResponse from pydantic import BaseModel from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline # ── Logging ──────────────────────────────────────────────────────────────────── logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # ── Config ───────────────────────────────────────────────────────────────────── MODEL_ID = "google/gemma-3-1b-it" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32 logger.info(f"Using device: {DEVICE} | dtype: {DTYPE}") # ── Global model state ───────────────────────────────────────────────────────── model_pipeline = None def load_model(): global model_pipeline logger.info(f"Loading model: {MODEL_ID} ...") hf_token = os.environ.get("HF_TOKEN") # Set this secret in HF Spaces settings tokenizer = AutoTokenizer.from_pretrained( MODEL_ID, token=hf_token, ) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=DTYPE, device_map="auto" if DEVICE == "cuda" else None, token=hf_token, ) if DEVICE == "cpu": model = model.to(DEVICE) model_pipeline = pipeline( "text-generation", model=model, tokenizer=tokenizer, device=0 if DEVICE == "cuda" else -1, ) logger.info("Model loaded successfully!") # ── Lifespan (startup / shutdown) ────────────────────────────────────────────── @asynccontextmanager async def lifespan(app: FastAPI): load_model() yield logger.info("Shutting down...") # ── FastAPI app ──────────────────────────────────────────────────────────────── app = FastAPI( title="Gemma-3-1B-IT API", description="FastAPI inference server for google/gemma-3-1b-it on HuggingFace Spaces", version="1.0.0", lifespan=lifespan, ) # ── Schemas ──────────────────────────────────────────────────────────────────── class Message(BaseModel): role: str # "user" or "assistant" content: str class ChatRequest(BaseModel): messages: List[Message] max_new_tokens: Optional[int] = 512 temperature: Optional[float] = 0.7 top_p: Optional[float] = 0.9 do_sample: Optional[bool] = True class GenerateRequest(BaseModel): prompt: str max_new_tokens: Optional[int] = 512 temperature: Optional[float] = 0.7 top_p: Optional[float] = 0.9 do_sample: Optional[bool] = True class ChatResponse(BaseModel): response: str model: str # ── Routes ───────────────────────────────────────────────────────────────────── @app.get("/", response_class=HTMLResponse) def root(): """Simple HTML landing page.""" return """ Gemma-3-1B-IT API

🤖 Gemma-3-1B-IT Inference API

Endpoints

Quick Example

curl -X POST /chat \\
  -H "Content-Type: application/json" \\
  -d '{
    "messages": [{"role": "user", "content": "Hello! Who are you?"}],
    "max_new_tokens": 200
  }'
""" @app.get("/health") def health(): """Health check endpoint.""" return { "status": "ok", "model": MODEL_ID, "device": DEVICE, "model_loaded": model_pipeline is not None, } @app.post("/chat", response_model=ChatResponse) def chat(request: ChatRequest): """ Chat endpoint — accepts a list of messages and returns the assistant reply. Uses the model's chat template automatically. """ if model_pipeline is None: raise HTTPException(status_code=503, detail="Model not loaded yet. Please retry.") try: # Build chat messages list messages = [{"role": m.role, "content": m.content} for m in request.messages] # Apply the chat template via the tokenizer tokenizer = model_pipeline.tokenizer prompt = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, ) outputs = model_pipeline( prompt, max_new_tokens=request.max_new_tokens, temperature=request.temperature, top_p=request.top_p, do_sample=request.do_sample, return_full_text=False, # return only the new tokens ) reply = outputs[0]["generated_text"].strip() return ChatResponse(response=reply, model=MODEL_ID) except Exception as e: logger.error(f"Chat error: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.post("/generate", response_model=ChatResponse) def generate(request: GenerateRequest): """ Raw text generation endpoint — pass a plain prompt string. """ if model_pipeline is None: raise HTTPException(status_code=503, detail="Model not loaded yet. Please retry.") try: outputs = model_pipeline( request.prompt, max_new_tokens=request.max_new_tokens, temperature=request.temperature, top_p=request.top_p, do_sample=request.do_sample, return_full_text=False, ) reply = outputs[0]["generated_text"].strip() return ChatResponse(response=reply, model=MODEL_ID) except Exception as e: logger.error(f"Generate error: {e}") raise HTTPException(status_code=500, detail=str(e))