import gradio as gr
import torch
import librosa
import numpy as np
from transformers import WhisperProcessor, WhisperForConditionalGeneration

print("Loading Whisper Tiny...")
whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-tiny")
whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny")
print("Whisper loaded!")

def transcribe_audio(audio_path):
    if audio_path is None:
        return "No audio recorded."
    try:
        speech, sr = librosa.load(audio_path, sr=16000)
        
        if np.abs(speech).max() > 0:
            speech = speech / np.abs(speech).max()

        whisper_inputs = whisper_processor(speech, sampling_rate=16000, return_tensors="pt")
        
        with torch.no_grad():
            generated_ids = whisper_model.generate(
                whisper_inputs.input_features,
                max_new_tokens=256,
                repetition_penalty=1.2,
                no_repeat_ngram_size=3
            )
        
        return whisper_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

    except Exception as e:
        return f"Error: {str(e)}"

demo = gr.Interface(
    fn=transcribe_audio,
    inputs=gr.Audio(type="filepath", label="Record Audio"),
    outputs=gr.Textbox(label="Transcription"),
    title="Medical Speech Recognition",
    description="Supports English and Arabic",
    flagging_mode="never"
)

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