src/mini_transformer/apps/server.py

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}