Spaces:
Sleeping
Sleeping
File size: 1,518 Bytes
2d52135 | 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 | from __future__ import annotations
import os
from fastapi import FastAPI, HTTPException
from omegaconf import DictConfig
from mini_transformer.inference import run_inference
from mini_transformer.model_loader import (
CONFIG_DIR_ENV,
CONFIG_NAME_ENV,
MODEL_NAME_ENV,
compose_config_from_dir,
compose_model_config,
ensure_models_root,
list_model_names,
)
app = FastAPI(title="Mini-Transformer Inference API")
@app.get("/healthz")
def healthz():
return {"status": "ok"}
@app.get("/generate")
def generate(text: str | None = None):
config_dir = os.environ.get(CONFIG_DIR_ENV)
config_name = os.environ.get(CONFIG_NAME_ENV, "config_inference")
if config_dir:
cfg: DictConfig = compose_config_from_dir(config_dir, config_name=config_name)
else:
ensure_models_root()
model_name = os.environ.get(MODEL_NAME_ENV)
if not model_name:
names = list_model_names()
if not names:
raise HTTPException(
status_code=503,
detail=(
"No models available under `trained_models/`. "
"Start the server with --model or place a model folder."
),
)
model_name = names[0]
cfg = compose_model_config(model_name, config_name=config_name)
if text is not None and text.strip():
cfg.input_text = text
outputs = run_inference(cfg)
return {"outputs": outputs}
|