| 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.basicConfig(level=logging.INFO) |
| logger = logging.getLogger(__name__) |
|
|
| |
| 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}") |
|
|
| |
| model_pipeline = None |
|
|
|
|
| def load_model(): |
| global model_pipeline |
| logger.info(f"Loading model: {MODEL_ID} ...") |
|
|
| hf_token = os.environ.get("HF_TOKEN") |
|
|
| 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!") |
|
|
|
|
| |
| @asynccontextmanager |
| async def lifespan(app: FastAPI): |
| load_model() |
| yield |
| logger.info("Shutting down...") |
|
|
|
|
| |
| 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, |
| ) |
|
|
|
|
| |
| class Message(BaseModel): |
| role: str |
| 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 |
|
|
|
|
| |
|
|
| @app.get("/", response_class=HTMLResponse) |
| def root(): |
| """Simple HTML landing page.""" |
| return """ |
| <!DOCTYPE html> |
| <html> |
| <head> |
| <title>Gemma-3-1B-IT API</title> |
| <style> |
| body { font-family: Arial, sans-serif; max-width: 800px; margin: 60px auto; padding: 0 20px; background: #f5f5f5; } |
| h1 { color: #333; } |
| a { color: #007bff; } |
| pre { background: #222; color: #0f0; padding: 15px; border-radius: 8px; overflow-x: auto; } |
| .card{ background: white; padding: 20px; border-radius: 10px; margin: 20px 0; box-shadow: 0 2px 8px rgba(0,0,0,0.1); } |
| </style> |
| </head> |
| <body> |
| <h1>🤖 Gemma-3-1B-IT Inference API</h1> |
| <div class="card"> |
| <h2>Endpoints</h2> |
| <ul> |
| <li><a href="/docs">📚 Interactive API Docs (Swagger UI)</a></li> |
| <li><a href="/health"><code>GET /health</code></a> — Health check</li> |
| <li><code>POST /chat</code> — Chat with message history</li> |
| <li><code>POST /generate</code> — Raw text generation</li> |
| </ul> |
| </div> |
| <div class="card"> |
| <h2>Quick Example</h2> |
| <pre>curl -X POST /chat \\ |
| -H "Content-Type: application/json" \\ |
| -d '{ |
| "messages": [{"role": "user", "content": "Hello! Who are you?"}], |
| "max_new_tokens": 200 |
| }'</pre> |
| </div> |
| </body> |
| </html> |
| """ |
|
|
|
|
| @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: |
| |
| messages = [{"role": m.role, "content": m.content} for m in request.messages] |
|
|
| |
| 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, |
| ) |
|
|
| 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)) |
|
|
