108 lines
3.3 KiB
Python
108 lines
3.3 KiB
Python
import os
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from fastapi import Depends, FastAPI, File, Header, HTTPException, UploadFile
|
|
from faster_whisper import WhisperModel
|
|
|
|
|
|
APP_NAME = "Sport Era Whisper ASR"
|
|
ALLOWED_EXTENSIONS = {"mp3", "m4a", "aac", "wav", "amr", "mp4", "3gp"}
|
|
|
|
app = FastAPI(title=APP_NAME)
|
|
model: Optional[WhisperModel] = None
|
|
|
|
|
|
def env_int(name: str, default: int) -> int:
|
|
try:
|
|
return int(os.getenv(name, str(default)))
|
|
except ValueError:
|
|
return default
|
|
|
|
|
|
def get_model() -> WhisperModel:
|
|
global model
|
|
if model is None:
|
|
model_name = os.getenv("ASR_MODEL", "small")
|
|
device = os.getenv("ASR_DEVICE", "cpu")
|
|
compute_type = os.getenv("ASR_COMPUTE_TYPE", "int8")
|
|
cache_dir = os.getenv("ASR_MODEL_CACHE", "/models")
|
|
model = WhisperModel(
|
|
model_name,
|
|
device=device,
|
|
compute_type=compute_type,
|
|
download_root=cache_dir,
|
|
)
|
|
return model
|
|
|
|
|
|
def verify_token(authorization: str = Header(default="")) -> None:
|
|
token = os.getenv("ASR_API_TOKEN", "").strip()
|
|
if not token:
|
|
return
|
|
expected = f"Bearer {token}"
|
|
if authorization != expected:
|
|
raise HTTPException(status_code=401, detail="Unauthorized")
|
|
|
|
|
|
def file_extension(filename: str) -> str:
|
|
return Path(filename or "").suffix.lower().lstrip(".")
|
|
|
|
|
|
@app.get("/health")
|
|
def health() -> dict:
|
|
return {
|
|
"status": "ok",
|
|
"service": APP_NAME,
|
|
"model": os.getenv("ASR_MODEL", "small"),
|
|
"device": os.getenv("ASR_DEVICE", "cpu"),
|
|
"compute_type": os.getenv("ASR_COMPUTE_TYPE", "int8"),
|
|
}
|
|
|
|
|
|
@app.post("/v1/transcribe")
|
|
async def transcribe(file: UploadFile = File(...), _: None = Depends(verify_token)) -> dict:
|
|
extension = file_extension(file.filename)
|
|
if extension not in ALLOWED_EXTENSIONS:
|
|
raise HTTPException(status_code=400, detail="Unsupported audio format")
|
|
|
|
max_size_mb = max(1, env_int("ASR_MAX_FILE_MB", 10))
|
|
content = await file.read(max_size_mb * 1024 * 1024 + 1)
|
|
if not content:
|
|
raise HTTPException(status_code=400, detail="Empty audio file")
|
|
if len(content) > max_size_mb * 1024 * 1024:
|
|
raise HTTPException(status_code=413, detail=f"Audio file is larger than {max_size_mb}MB")
|
|
|
|
temp_path = ""
|
|
try:
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix=f".{extension}") as temp_file:
|
|
temp_file.write(content)
|
|
temp_path = temp_file.name
|
|
|
|
language = os.getenv("ASR_LANGUAGE", "").strip() or None
|
|
beam_size = max(1, env_int("ASR_BEAM_SIZE", 5))
|
|
start_time = time.perf_counter()
|
|
segments, info = get_model().transcribe(
|
|
temp_path,
|
|
language=language,
|
|
beam_size=beam_size,
|
|
vad_filter=True,
|
|
)
|
|
text = "".join(segment.text for segment in segments).strip()
|
|
duration_ms = int((time.perf_counter() - start_time) * 1000)
|
|
return {
|
|
"text": text,
|
|
"language": info.language,
|
|
"language_probability": info.language_probability,
|
|
"duration": info.duration,
|
|
"duration_ms": duration_ms,
|
|
}
|
|
finally:
|
|
if temp_path:
|
|
try:
|
|
os.unlink(temp_path)
|
|
except OSError:
|
|
pass
|