handler.py

import runpod
import base64
import tempfile
import requests
import os
import enum
import sys
import traceback
import logging
import subprocess
import soundfile as sf
from runpod import RunPodLogger
from litellm import completion
from litellm.exceptions import Timeout
from dotenv import load_dotenv

# Ensure we can import MegaASR from the src folder
sys.path.append(os.path.abspath("src"))
from MegaASR.model.megaASR import MegaASR

load_dotenv(verbose=True, override=True)

logger = RunPodLogger()

DEBUG = os.getenv("DEBUG", "false").lower() in ("true", "1", "yes")
logger.info(f"Debug mode: {DEBUG}")

if DEBUG:
    logging.getLogger("MegaASR").setLevel(logging.DEBUG)

# Initialize vLLM options from env or use small GPU defaults
vllm_kwargs = {
    "gpu_memory_utilization": float(os.getenv("GPU_MEMORY_UTILIZATION", "0.85")),
    "max_model_len": int(os.getenv("MAX_MODEL_LEN", "8192")),
    "max_num_seqs": int(os.getenv("MAX_NUM_SEQS", "1")),
    "max_num_batched_tokens": int(os.getenv("MAX_NUM_BATCHED_TOKENS", "2048")),
}
logger.info(f"vLLM config: {vllm_kwargs}")

# Initialize model at startup
logger.info("Initializing MegaASR model...")
model = MegaASR(
    model_path="ckpt/Mega-ASR/Qwen3-ASR-1.7B",
    lora_dir="ckpt/Mega-ASR/mega-asr-merged",
    routing_enabled=False,
    backend="vllm",
    vllm_apply_lora_on_load=True,
    vllm_materialized_lora_dir="ckpt/Mega-ASR/mega-asr-vllm-materialized",
    **vllm_kwargs,
)
logger.info("MegaASR model initialized successfully.")

USE_LITELLM = False
LITELLM_MODEL = os.getenv("LITELLM_MODEL")
LITELLM_API_KEY = os.getenv("LITELLM_API_KEY")
LITELLM_API_VERSION = os.getenv("LITELLM_API_VERSION")
LITELLM_API_BASE = os.getenv("LITELLM_API_BASE")

if LITELLM_MODEL and LITELLM_API_KEY:
    USE_LITELLM = True
    logger.info("Using LiteLLM for translation and hallucination detection")


def base64_to_tempfile(base64_data):
    logger.debug("Decoding base64 audio data to tempfile.")
    try:
        audio_data = base64.b64decode(base64_data)
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
        with open(temp_file.name, "wb") as file:
            file.write(audio_data)
        logger.debug(f"Base64 audio written to tempfile: {temp_file.name}")
        return temp_file.name
    except Exception as e:
        logger.error(f"Error in base64_to_tempfile: {e}")
        raise


def download_url_to_mp3(url):
    logger.debug(f"Downloading audio from URL: {url}")
    try:
        response = requests.get(url)
        logger.debug(f"Download response status: {response.status_code}")
        if response.status_code != 200:
            logger.error(f"Failed to download file from URL: {url}")
            raise Exception("Failed to download file from URL")
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
        temp_file.write(response.content)
        temp_file.close()
        logger.debug(f"Audio downloaded and saved to tempfile: {temp_file.name}")
        return temp_file.name
    except Exception as e:
        logger.error(f"Error in download_url_to_mp3: {e}")
        raise


def convert_to_wav(input_path):
    """
    Normalizes input audio format to 16kHz mono WAV using ffmpeg.
    This guarantees compatibility with soundfile.
    """
    logger.debug(f"Normalizing audio {input_path} to WAV using ffmpeg.")
    temp_wav = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    temp_wav.close()
    
    cmd = [
        "ffmpeg", "-y",
        "-i", input_path,
        "-ar", "16000",
        "-ac", "1",
        "-c:a", "pcm_s16le",
        temp_wav.name
    ]
    try:
        subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        logger.debug(f"Audio normalization successful: {temp_wav.name}")
        return temp_wav.name
    except Exception as e:
        logger.error(f"FFmpeg normalization failed: {e}")
        if os.path.exists(temp_wav.name):
            os.remove(temp_wav.name)
        raise


class ExitCode(enum.Enum):
    SUCCESS = 0
    ERROR_TRANSLATING = 1
    TIMEOUT = 2


TRANSLATION_PROMPT = (
    """You are an expert multilingual translation assistant. Your task is to produce a single, coherent text in the specified target language, ensuring the final output is natural and grammatically correct.

**Your Process:**
1.  **Identify:** Pinpoint the phrases in the source text that are NOT in the target language.
2.  **Translate:** Translate ONLY those identified parts into the target language.
3.  **Preserve:** Leave any text that is ALREADY in the target language completely unchanged.
4.  **Integrate:** Combine the newly translated parts and the preserved parts into a single, seamless output text.

**CRITICAL GUIDELINES:**
- **Hotwords:** The user has provided a list of "hotwords". You MUST preserve these words exactly as they appear, including their original capitalization. DO NOT translate them.
- **Output:** Output ONLY the final integrated text. Do not add commentary, explanations, or any other text.

---
**Example 1: Mixed Language with Hotwords**
- **Target Language:** `nl`
- **Hotwords:** ["RunPod", "WhisperX"]
- **Source Text:** `I am running the new WhisperX model on the RunPod server. Ik deel nu mijn scherm.`
- **Your Correct Output:** `Ik draai het nieuwe WhisperX model op de RunPod server. Ik deel nu mijn scherm.`
---
**Example 2: Simple Translation with Hotwords**
- **Target Language:** `es`
- **Hotwords:** ["Project Phoenix"]
- **Source Text:** `Let's discuss the status of Project Phoenix.`
- **Your Correct Output:** `Discutamos el estado de Project Phoenix.`
---
"""
)


def translate_text(
    text: str,
    language: str,
    hotwords: str = None,
) -> tuple[str, int]:
    logger.debug(f"Translating text to '{language}' with hotwords: {hotwords}. Text: {text[:100]}...")

    user_prompt_parts = [
        f"Target Language: `{language}`"
    ]

    if hotwords:
        user_prompt_parts.append(f"Hotwords: [{hotwords}]")

    user_prompt_parts.append(f'\nSource Text: "{text}"')

    user_prompt = "\n".join(user_prompt_parts)

    try:
        response = completion(
            model=str(LITELLM_MODEL),
            messages=[
                {"role": "system", "content": TRANSLATION_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            api_key=LITELLM_API_KEY,
            api_version=LITELLM_API_VERSION,
            api_base=LITELLM_API_BASE,
        )
        translated_text = response.choices[0].message.content
        logger.debug(f"Translation response: {translated_text}")
        return translated_text, ExitCode.SUCCESS.value

    except Timeout as e:
        logger.error(f"Timeout in translate_text: {e}")
        return text, ExitCode.TIMEOUT.value
    except Exception as e:
        logger.error(f"Error in translate_text: {e}")
        return text, ExitCode.ERROR_TRANSLATING.value


HALLUCINATION_PROMPT = """You are an expert QA analyst specializing in reviewing speech-to-text (ASR) transcripts. Your task is to identify "hallucinations" without access to the original audio.

A hallucination is defined as text that is highly unlikely to have been spoken by a coherent human. You must infer this from textual evidence alone.

**Crucially, distinguish between:**
- **Hallucinations (High Penalty):** Invented phrases, nonsensical "word salad," illogical topic shifts, or repetitive loops that suggest ASR model failure.
- **Minor Inaccuracies (Low/No Penalty):** Simple transcription errors, misheard names, or natural filler words (e.g., "um," "uh"). Do not score these as severe hallucinations.

**Evaluation Criteria (assign ONE score based on the MOST SEVERE issue found):**
- **0.0 - No Hallucination:** The text is coherent and sounds like natural human speech.
- **0.1-0.3 - Minor:** Mostly coherent but contains a few slightly awkward or out-of-place words.
- **0.4-0.5 - Moderate:** Contains distracting nonsensical phrases or confusing sentences that disrupt the flow.
- **0.6-0.9 - Severe:** Multiple incoherent passages, significant logical breaks, or repetitive loops that obscure meaning.
- **1.0 - Complete:** The text is almost entirely gibberish, stuck in a repetitive loop, OR consists solely of a list of the provided "Hotwords" with no surrounding conversational context.

**Special Case: Isolated Keyword (Hotword) Lists**
If the ASR output consists ONLY of a list of the provided "Hotwords" (e.g., just the names themselves, separated by commas or spaces) and nothing else, this is a sign of complete ASR failure. **This case MUST be scored as 1.0.**

**Primary Hallucination Signals to Look For:**
1.  **Logical Incoherence or Nonsense:** Grammatically malformed or self-contradictory sentences.
2.  **Semantic Repetition or Looping:** The same idea or phrase repeated unnaturally.
3.  **Abrupt and Illogical Topic Shifts:** Sudden changes in topic without logical transition.
4.  **Inappropriate Jargon or "Word Salad":** Technical terms used completely out of context.

Respond ONLY with a valid JSON object in the following format:
{
  "hallucination_score": <score from 0.0 to 1.0>,
  "reason": "<A concise, max 20-word explanation for your score, citing the most significant issue.>"
}
"""


def detect_hallucination(text: str, hotwords: str = "") -> tuple[float, str]:
    logger.debug(f"Checking for hallucinations in text: {text[:100]}...")

    if not (LITELLM_MODEL and LITELLM_API_KEY):
        return 0.0, ""

    try:
        hotwords_info = f"Hotwords: {hotwords}\n" if hotwords else ""
        user_content = (
            f"Analyze this ASR transcript for signs of hallucination.\n"
            f"{hotwords_info}\n"
            f"Transcript: {text}"
        )

        response = completion(
            model=str(LITELLM_MODEL),
            messages=[
                {
                    "role": "system",
                    "content": HALLUCINATION_PROMPT,
                },
                {
                    "role": "user",
                    "content": user_content,
                },
            ],
            api_key=LITELLM_API_KEY,
            api_version=LITELLM_API_VERSION,
            api_base=LITELLM_API_BASE,
            response_format={"type": "json_object"},
        )
        result = response.choices[0].message.content
        import json
        parsed = json.loads(result)
        score = parsed.get("hallucination_score", 0.0)
        try:
            score = float(score)
        except (TypeError, ValueError):
            score = 0.0
        return score, parsed.get("reason", "")
    except Exception as e:
        logger.error(f"Error in hallucination detection: {e}")
        return 0.0, f"Detection error: {str(e)}"


def parse_segments(raw_results, audio_duration):
    """
    Parses segment/word information dynamically from Qwen3-ASR vLLM output.
    Uses fallback if no structured object attributes exist.
    """
    segments = []
    if isinstance(raw_results, list):
        for item in raw_results:
            text = getattr(item, "text", str(item)).strip()
            start = getattr(item, "start", 0.0)
            end = getattr(item, "end", audio_duration)
            
            words = None
            if hasattr(item, "words") and item.words:
                words = []
                for w in item.words:
                    words.append({
                        "word": getattr(w, "word", str(w)),
                        "start": getattr(w, "start", 0.0),
                        "end": getattr(w, "end", 0.0)
                    })
            
            segments.append({
                "text": text,
                "start": start,
                "end": end,
                "words": words
            })
    else:
        segments.append({
            "text": str(raw_results).strip(),
            "start": 0.0,
            "end": audio_duration,
            "words": None
        })
    return segments


def clean_up_audio(*paths):
    for path in paths:
        try:
            if path and os.path.exists(path):
                if not DEBUG:
                    logger.debug(f"Removing temp file: {path}")
                    os.remove(path)
                else:
                    logger.debug(f"Not removing temp file (DEBUG mode): {path}")
        except Exception as e:
            logger.error(f"Error in clean_up_audio for {path}: {e}")


def handler(event):
    logger.debug(f"Handler called with event: {event}")
    raw_audio_path = None
    wav_audio_path = None

    try:
        # ---------------------- Parse Input ----------------------
        job_input = event["input"]
        job_input_audio_base_64 = job_input.get("audio_base_64")
        job_input_audio_url = job_input.get("audio")
        job_input_language = job_input.get("language", None)
        hotwords = job_input.get("hotwords", None)
        enable_timestamps = job_input.get("enable_timestamps", False)
        disable_hallucination_detection = job_input.get(
            "disable_hallucination_detection", False
        )
        disable_translation = job_input.get("disable_translation", False)

        # Metadata passthrough
        conversation_id = job_input.get("conversation_id", "")
        conversation_chunk_id = job_input.get("conversation_chunk_id", "")
        metadata_str = job_input.get("metadata_str", "")

        logger.info(f"Job input: {job_input}")

        # ---------------------- Audio Preparation ----------------------
        if job_input_audio_base_64:
            logger.debug("Audio input provided as base64.")
            raw_audio_path = base64_to_tempfile(job_input_audio_base_64)
        elif job_input_audio_url:
            if job_input_audio_url.startswith("http"):
                logger.debug("Audio input provided as URL.")
                raw_audio_path = download_url_to_mp3(job_input_audio_url)
            else:
                logger.debug("Audio input provided as local file.")
                if not os.path.exists(job_input_audio_url):
                    raise ValueError(f"Local file {job_input_audio_url} does not exist.")
                raw_audio_path = job_input_audio_url
        else:
            raise ValueError("No audio input provided")

        # Convert/Normalize audio to 16kHz mono WAV for soundfile and Mega-ASR compatibility
        wav_audio_path = convert_to_wav(raw_audio_path)
        
        # Get audio duration
        audio_info = sf.info(wav_audio_path)
        audio_duration = audio_info.duration
        logger.debug(f"Audio file duration: {audio_duration:.2f}s")

        # ---------------------- Transcription ----------------------
        logger.debug("Running Mega-ASR transcription.")
        
        transcribe_kwargs = {}
        if enable_timestamps:
            transcribe_kwargs["return_objects"] = True
            # Check if we should pass return_time_stamps to the underlying qwen_asr model
            transcribe_kwargs["return_time_stamps"] = True

        # Run model inference
        raw_result = model.infer(wav_audio_path, **transcribe_kwargs)

        if enable_timestamps:
            # Parse segments dynamically from the returned objects
            generated_segments = parse_segments(raw_result, audio_duration)
            joined_text = " ".join([segment["text"] for segment in generated_segments]).strip()
        else:
            joined_text = str(raw_result).strip()
            generated_segments = []

        logger.debug(f"Transcription complete: {joined_text[:200]}...")

        # ---------------------- Translation ----------------------
        translation_text = None
        translation_error = None

        # For Mega-ASR, we bypass language detection comparison and rely on the
        # translation model prompt to translate only non-matching parts, as
        # Mega-ASR/Qwen3-ASR doesn't reliably expose a fast language confidence score.
        needs_translation = (
            not disable_translation and
            job_input_language is not None
        )

        if needs_translation:
            logger.debug(f"Translating text to '{job_input_language}' with hotwords: {hotwords}.")
            translation_text, translation_error = translate_text(
                joined_text,
                job_input_language,
                hotwords
            )
        else:
            logger.debug("Translation not required or disabled.")

        # ---------------------- Hallucination Detection ----------------------
        hallucination_score = None
        hallucination_reason = None

        if USE_LITELLM and not disable_hallucination_detection:
            text_to_analyze_for_hallucination = translation_text if translation_text else joined_text
            if text_to_analyze_for_hallucination:
                hallucination_score, hallucination_reason = detect_hallucination(
                    text_to_analyze_for_hallucination,
                    hotwords
                )
                if hallucination_score >= 0.9:
                    logger.info(f"Severe hallucination detected (score={hallucination_score}): {hallucination_reason}")
                elif hallucination_score >= 0.5:
                    logger.info(f"Moderate hallucination detected (score={hallucination_score}): {hallucination_reason}")

        # Build response payload matching runpod-whisper interface exactly
        result = {
            "conversation_id": conversation_id,
            "conversation_chunk_id": conversation_chunk_id,
            "metadata_str": metadata_str,
            "enable_timestamps": enable_timestamps,
            "language": job_input_language,
            "detected_language": "en" if "zh" not in joined_text else "zh",  # Simple heuristic fallback
            "detected_language_confidence": 1.0,
            "translation_text": translation_text,
            "translation_error": translation_error,
            "hallucination_score": hallucination_score,
            "hallucination_reason": hallucination_reason if hallucination_score is not None else "",
            "joined_text": joined_text,
        }

        if enable_timestamps:
            result["segments"] = generated_segments

        # Clean up files
        clean_up_audio(raw_audio_path, wav_audio_path)
        return result

    except Exception as e:
        logger.error(f"Unhandled error: {str(e)}")
        logger.error(traceback.format_exc())

        clean_up_audio(raw_audio_path, wav_audio_path)

        common_meta = {
            "conversation_id": locals().get("conversation_id", ""),
            "conversation_chunk_id": locals().get("conversation_chunk_id", ""),
            "metadata_str": locals().get("metadata_str", ""),
            "enable_timestamps": locals().get("enable_timestamps", False),
            "language": locals().get("job_input_language", None),
        }

        return {
            **common_meta,
            "error": str(e),
            "message": "An unhandled error occurred while processing the request.",
        }


runpod.serverless.start({"handler": handler})