"""MedASR Transcription API — HuggingFace Space.

Loads google/medasr (105M Conformer model) and exposes a Gradio
interface + automatic REST API for remote transcription.

Uses AutoModelForCTC + AutoProcessor directly instead of the pipeline
API to avoid the _torch_extract_fbank_features bug in transformers 5.x.
"""

import os
import gradio as gr
import torch
import librosa
from transformers import AutoModelForCTC, AutoProcessor

HF_TOKEN = os.getenv("HF_TOKEN")

MODEL_ID = "google/medasr"
DEVICE = "cpu"

print("[MedASR-Space] Loading model …")
processor = AutoProcessor.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = AutoModelForCTC.from_pretrained(MODEL_ID, token=HF_TOKEN).to(DEVICE)
print("[MedASR-Space] Model ready.")


def transcribe(audio_path: str) -> str:
    """Transcribe audio file using MedASR."""
    if audio_path is None:
        return ""

    speech, sample_rate = librosa.load(audio_path, sr=16000)
    inputs = processor(speech, sampling_rate=sample_rate, return_tensors="pt", padding=True)
    inputs = inputs.to(DEVICE)

    with torch.no_grad():
        outputs = model.generate(**inputs)

    text = processor.batch_decode(outputs)[0]
    return text.replace("</s>", "").replace("<s>", "").strip()


demo = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(type="filepath", label="Audio"),
    outputs=gr.Textbox(label="Transcription"),
    title="MedASR — Medical Speech Recognition",
    description="Google MedASR (Conformer 105M) for medical dictation and EMS transcription.",
)

demo.launch(show_error=True)