Spaces:
Running
Running
File size: 3,131 Bytes
c1de90b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | from __future__ import annotations
import os
from threading import Lock
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from .inference import generate_code, load_generator
class GenerateRequest(BaseModel):
instruction: str = Field(..., min_length=1)
input: str = ""
max_new_tokens: int = Field(default=512, ge=1, le=4096)
temperature: float = Field(default=0.2, ge=0.0, le=2.0)
top_p: float = Field(default=0.95, ge=0.01, le=1.0)
safety: bool = True
class GenerateResponse(BaseModel):
completion: str
time_taken: float = 0.0
app = FastAPI(title="Gemma Code LLM API", version="0.1.0")
_model_state = None
_model_lock = Lock()
def _load_state():
global _model_state
if _model_state is not None:
return _model_state
with _model_lock:
if _model_state is not None:
return _model_state
base_model = os.getenv("BASE_MODEL", "google/gemma-3-1b-it")
adapter = os.getenv("ADAPTER_PATH") or None
quantization = os.getenv("QUANTIZATION", "none")
dtype = os.getenv("DTYPE", "auto")
trust_remote_code = os.getenv("TRUST_REMOTE_CODE", "0") == "1"
_model_state = load_generator(
base_model,
adapter=adapter,
quantization=quantization,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
return _model_state
@app.get("/health")
def health() -> dict[str, str]:
return {"status": "ok"}
@app.post("/generate", response_model=GenerateResponse)
def generate(request: GenerateRequest) -> GenerateResponse:
import time
start_time = time.time()
try:
model, tokenizer, torch = _load_state()
completion = generate_code(
model,
tokenizer,
torch,
instruction=request.instruction,
input_text=request.input,
max_new_tokens=request.max_new_tokens,
temperature=request.temperature,
top_p=request.top_p,
safety=request.safety,
)
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
except Exception as exc:
raise HTTPException(status_code=500, detail=str(exc)) from exc
elapsed_time = time.time() - start_time
return GenerateResponse(completion=completion, time_taken=elapsed_time)
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
@app.get("/", response_class=HTMLResponse)
def read_root():
static_file_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "static", "index.html")
try:
with open(static_file_path, "r", encoding="utf-8") as f:
return HTMLResponse(content=f.read())
except FileNotFoundError:
return HTMLResponse(content="<h1>UI Files not found</h1><p>Ensure that 'static' folder exists in the project root.</p>", status_code=404)
static_dir_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "static")
app.mount("/static", StaticFiles(directory=static_dir_path), name="static")
|