from __future__ import annotations

from typing import Any
from asg_transformer import ASGTransformer

_model: ASGTransformer | None = None


def load_model(model_dir: str = ".") -> ASGTransformer:
    global _model
    if _model is None:
        _model = ASGTransformer.from_pretrained(model_dir)
    return _model


def predict(inputs: dict[str, Any]) -> dict[str, Any]:
    model = load_model()
    text = inputs.get("text") or inputs.get("inputs")
    if not isinstance(text, str) or not text.strip():
        raise ValueError("A non-empty 'text' or 'inputs' field is required")
    return model.generate(
        text,
        max_steps=inputs.get("max_steps"),
        beam_width=inputs.get("beam_width"),
        transition_weight=inputs.get("transition_weight"),
        total_duration_minutes=inputs.get("total_duration_minutes"),
        language=inputs.get("language", "en"),
    ).to_dict()