app.py

import gradio as gr
from transformers import AutoModelForCTC, AutoFeatureExtractor, AutoTokenizer
import torch
import numpy as np
import warnings
import librosa

warnings.filterwarnings("ignore")

MODEL_ID = "google/medasr"

model = None
feature_extractor = None
tokenizer = None

def normalize_audio(audio):
    """RMS归一化"""
    rms = np.sqrt(np.mean(audio ** 2))
    if rms > 0:
        audio = audio / rms
        audio = np.clip(audio, -1.0, 1.0)
    return audio

def remove_silence(audio, sample_rate, threshold=0.01):
    """去除静音段"""
    energy = np.abs(audio)
    above_threshold = energy > threshold
    if not np.any(above_threshold):
        return audio
    start = np.where(above_threshold)[0][0]
    end = np.where(above_threshold)[0][-1]
    buffer = int(0.1 * sample_rate)
    start = max(0, start - buffer)
    end = min(len(audio), end + buffer)
    return audio[start:end]

def load_model_with_token(hf_token):
    global model, feature_extractor, tokenizer

    if not hf_token or not hf_token.strip():
        return gr.update(interactive=False, value="❌ Token cannot be empty!"), gr.update(interactive=False)

    try:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print("🔄 Loading model components...")
        print(f"📱 Device: {device}")

        feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_ID, token=hf_token.strip())
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=hf_token.strip())
        model = AutoModelForCTC.from_pretrained(MODEL_ID, token=hf_token.strip()).to(device)
        model.eval()

        print(f"✅ Loaded: {type(feature_extractor)}, {type(tokenizer)}, {type(model)}")
        
        return gr.update(interactive=False, value="✅ Model Loaded Successfully!"), gr.update(interactive=True)
    except Exception as e:
        print(f"Error loading model: {e}")
        import traceback
        traceback.print_exc()
        return gr.update(interactive=True, value=f"❌ Error: {str(e)}"), gr.update(interactive=False)


def transcribe_audio(audio_input):
    global model, feature_extractor, tokenizer

    if audio_input is None:
        return "⚠️ Please upload or record audio."

    if model is None or feature_extractor is None:
        return "❌ Please load the model first!"

    try:
        # 1. 解包音频
        sample_rate, waveform = audio_input

        # 2. 转单声道
        if waveform.ndim == 2:
            waveform = waveform[:, 0]

        # 3. 转换为 float32 并归一化
        if waveform.dtype == np.int16:
            waveform = waveform.astype(np.float32) / 32768.0
        elif waveform.dtype != np.float32:
            waveform = waveform.astype(np.float32)

        # 4. RMS归一化
        waveform = normalize_audio(waveform)

        # 5. 去除静音
        waveform = remove_silence(waveform, sample_rate)

        # 6. 检查长度
        duration = len(waveform) / sample_rate
        if duration < 0.1:
            return "⚠️ Audio is too short."
        if duration > 60:
            return "⚠️ Audio is too long."

        # 7. 重采样
        if sample_rate != 16000:
            waveform = librosa.resample(waveform, orig_sr=sample_rate, target_sr=16000)
            sample_rate = 16000

        # 8. 特征提取
        inputs = feature_extractor(
            waveform,
            sampling_rate=sample_rate,
            return_tensors="pt",
        )
        inputs = {k: v.to(model.device) for k, v in inputs.items()}

        # 【修复部分】
        # 自动查找包含特征数据的键（可能是 'input_features' 或 'input_values'）
        # 过滤掉 'attention_mask'，找到真正的输入 Tensor
        input_tensor = None
        for key, val in inputs.items():
            if isinstance(val, torch.Tensor) and val.ndim > 1:
                input_tensor = val
                break
        
        if input_tensor is None:
            return "❌ Error: Could not extract audio features."

        # 安全地获取 stride，如果不存在则默认为 4
        stride = 4
        if hasattr(feature_extractor, 'stride'):
            s = feature_extractor.stride
            stride = s[0] if isinstance(s, (list, tuple)) else s
        
        # 动态计算 max_length
        max_length = input_tensor.shape[1] // stride + 50

        # 9. Beam search 解码
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_length=max_length,
                num_beams=8,  # Beam search 提升准确率
                temperature=1.0,
            )

        # 10. 解码
        transcription = tokenizer.batch_decode(outputs.tolist(), skip_special_tokens=True)[0]

        # 11. 后处理
        transcription = transcription.strip()
        import re
        transcription = re.sub(r'\s+', ' ', transcription)

        return transcription if transcription else "⚠️ No speech detected."

    except Exception as e:
        import traceback
        traceback.print_exc()
        return f"❌ Transcription error: {str(e)}"


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🏥 MedASR - Medical Speech Recognition")
    gr.Markdown("Optimized for medical dictation with Beam Search decoding.")
    
    gr.Markdown("---")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## 🔑 Authentication")
            hf_token = gr.Textbox(label="HuggingFace Token", type="password", placeholder="hf_...")
            load_model_btn = gr.Button("🔑 Load Model", variant="primary", size="lg")
            
            gr.Markdown("## 📝 Tips")
            gr.Markdown("""
            - Speak **clearly and slowly**
            - Use **medical terms**
            - Short audio (3-10s) is best
            - Quiet environment
            """)

        with gr.Column(scale=2):
            gr.Markdown("## 🎙️ Input & Result")
            audio_input = gr.Audio(sources=["microphone", "upload"], type="numpy")
            
            with gr.Row():
                transcribe_btn = gr.Button("🔄 Transcribe", variant="secondary", size="lg", interactive=False)
                clear_btn = gr.Button("🗑️ Clear", variant="stop", size="lg")
            
            output_text = gr.Textbox(label="Result", lines=12, placeholder="...")
            audio_info = gr.Textbox(label="Info", lines=2, interactive=False)

    def transcribe_wrapper(audio_in):
        res = transcribe_audio(audio_in)
        info = f"Status: Success" if "❌" not in res and "⚠️" not in res else "Status: Check result"
        return res, info

    load_model_btn.click(
        fn=load_model_with_token,
        inputs=[hf_token],
        outputs=[load_model_btn, transcribe_btn]
    )

    transcribe_btn.click(
        fn=transcribe_wrapper,
        inputs=[audio_input],
        outputs=[output_text, audio_info]
    )
    
    clear_btn.click(
        fn=lambda: ("", "Ready"),
        outputs=[output_text, audio_info]
    )

if __name__ == "__main__":
    demo.launch()