Files
speech2text/docker/transcription-worker/main.py
2026-06-02 17:35:45 +08:00

126 lines
4.1 KiB
Python

"""
Speech2Text MBIP — Transcription Worker
FastAPI service wrapping faster-whisper for Bahasa Melayu transcription.
Internal-only service; not exposed to internet.
"""
import os
import tempfile
import logging
from pathlib import Path
from typing import Optional
from fastapi import FastAPI, UploadFile, File, HTTPException
from pydantic import BaseModel
from faster_whisper import WhisperModel
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI(title="Speech2Text Transcription Worker", docs_url=None, redoc_url=None)
# Config from environment
MODEL_SIZE = os.getenv("WHISPER_MODEL", "small")
LANGUAGE = os.getenv("WHISPER_LANGUAGE", "ms")
DEVICE = os.getenv("WHISPER_DEVICE", "cpu")
COMPUTE_TYPE = os.getenv("WHISPER_COMPUTE_TYPE", "int8")
logger.info(f"Loading Whisper model: {MODEL_SIZE}, device: {DEVICE}, compute: {COMPUTE_TYPE}")
# Load model once at startup
model = WhisperModel(MODEL_SIZE, device=DEVICE, compute_type=COMPUTE_TYPE)
logger.info("Whisper model loaded successfully.")
class TranscribeResponse(BaseModel):
success: bool
transcript: Optional[str] = None
language: Optional[str] = None
confidence: Optional[float] = None
duration_seconds: Optional[float] = None
error: Optional[str] = None
@app.get("/health")
def health_check():
return {"status": "ok", "model": MODEL_SIZE, "language": LANGUAGE}
@app.post("/transcribe", response_model=TranscribeResponse)
async def transcribe_audio(audio: UploadFile = File(...)):
"""
Terima fail audio, transkripkan ke teks Bahasa Melayu.
Hanya accessible dari dalam Docker network sahaja.
"""
# Validate content type
allowed_types = {
"audio/mpeg", "audio/mp3", "audio/wav", "audio/x-wav",
"audio/mp4", "audio/x-m4a", "audio/aac", "video/mp4",
"audio/ogg", "audio/flac", "audio/webm",
}
if audio.content_type and audio.content_type not in allowed_types:
logger.warning(f"Rejected content type: {audio.content_type}")
raise HTTPException(status_code=400, detail=f"Unsupported audio format: {audio.content_type}")
# Save to temp file
suffix = Path(audio.filename or "audio.mp3").suffix or ".mp3"
try:
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
content = await audio.read()
tmp.write(content)
tmp_path = tmp.name
logger.info(f"Transcribing: {audio.filename}, size: {len(content)} bytes")
# Run transcription
segments, info = model.transcribe(
tmp_path,
language=LANGUAGE,
beam_size=5,
vad_filter=True,
)
# Collect segments
all_text = []
total_confidence = []
for segment in segments:
all_text.append(segment.text.strip())
# faster-whisper segments have avg_logprob; convert to 0-1 range
if hasattr(segment, "avg_logprob") and segment.avg_logprob is not None:
import math
confidence = min(1.0, max(0.0, math.exp(segment.avg_logprob)))
total_confidence.append(confidence)
transcript = " ".join(all_text).strip()
avg_confidence = sum(total_confidence) / len(total_confidence) if total_confidence else None
duration = info.duration if hasattr(info, "duration") else None
logger.info(f"Transcription complete: {len(transcript)} chars, lang: {info.language}")
return TranscribeResponse(
success=True,
transcript=transcript,
language=info.language,
confidence=round(avg_confidence, 4) if avg_confidence else None,
duration_seconds=round(duration, 2) if duration else None,
)
except Exception as e:
logger.error(f"Transcription failed: {e}", exc_info=True)
return TranscribeResponse(
success=False,
error=str(e),
)
finally:
# Always cleanup temp file
if "tmp_path" in locals():
try:
os.unlink(tmp_path)
except Exception:
pass