publisher/tasks/transcribe.py

"""
BASYX V11 — Transcribe Task
Production Version
"""

from faster_whisper import WhisperModel
import tempfile
import os
import httpx
import asyncio

# --------------------------------------------------
# GLOBAL MODEL (LOAD ONCE)
# --------------------------------------------------

MODEL = WhisperModel(
    model_size_or_path="base",
    device="cpu",
    compute_type="int8"
)


# --------------------------------------------------
# HELPERS
# --------------------------------------------------

async def save_upload(file):
    """Save uploaded file to temp path"""
    suffix = os.path.splitext(file.filename)[-1]

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    tmp.write(await file.read())
    tmp.close()

    return tmp.name


async def download_url(url: str):
    """Download media from URL safely"""
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")

    async with httpx.AsyncClient(timeout=300) as client:
        async with client.stream("GET", url) as r:
            r.raise_for_status()
            async for chunk in r.aiter_bytes():
                tmp.write(chunk)

    tmp.close()
    return tmp.name


def build_segments(segments):
    """Normalize whisper segments"""
    results = []

    for seg in segments:
        results.append({
            "start": round(seg.start, 2),
            "end": round(seg.end, 2),
            "text": seg.text.strip()
        })

    return results


# --------------------------------------------------
# MAIN TASK ENTRYPOINT
# --------------------------------------------------

async def run(context):
    """
    Expected context:
        context.input_file
        context.url_input
    """

    media_path = None

    try:
        # ----------------------------
        # INPUT RESOLUTION
        # ----------------------------

        if context.input_file:
            media_path = await save_upload(context.input_file)

        elif context.url_input:
            media_path = await download_url(context.url_input)

        else:
            return {
                "status": "error",
                "message": "No file or URL provided"
            }

        # ----------------------------
        # TRANSCRIPTION
        # ----------------------------

        segments, info = await asyncio.to_thread(
            MODEL.transcribe,
            media_path,
            beam_size=5
        )

        segment_list = build_segments(segments)

        full_text = " ".join(s["text"] for s in segment_list)

        # ----------------------------
        # OUTPUT
        # ----------------------------

        return {
            "status": "success",
            "language": info.language,
            "duration": info.duration,
            "segments": segment_list,
            "text": full_text
        }

    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }

    finally:
        if media_path and os.path.exists(media_path):
            os.remove(media_path)