File size: 903 Bytes
b7fa3e3 | 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 | 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()
|