import gradio as gr
from transformers import WhisperProcessor, WhisperForConditionalGeneration, pipeline
import torch
import soundfile as sf

# Load Whisper model and processor from Hugging Face
model_name = "openai/whisper-large-v3"
processor = WhisperProcessor.from_pretrained(model_name)
model = WhisperForConditionalGeneration.from_pretrained(model_name)

# Ensure the model is using the correct device (GPU or CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Function to handle transcription with language set to English by default
def transcribe(audio_path):
    # Load audio from file
    audio, sampling_rate = sf.read(audio_path)

    # Process the audio to get input features
    input_features = processor(audio, sampling_rate=sampling_rate, return_tensors="pt").input_features.to(device)

    # Generate transcription with attention_mask and correct input_features
    attention_mask = torch.ones(input_features.shape, dtype=torch.long, device=device)
    generated_ids = model.generate(
        input_features=input_features,
        attention_mask=attention_mask,
        language="en"  # Force translation to English
    )
    
    # Decode transcription
    transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return transcription

# Create a Gradio Interface
interface = gr.Interface(
    fn=transcribe, 
    inputs=gr.Audio(sources="upload", type="filepath"), 
    outputs="text",
    title="Whisper Speech-to-Text API",
    description="Upload an audio file and get a transcription using OpenAI's Whisper model from Hugging Face."
)

# Launch the interface as an API
interface.launch()