import gradio as gr
import whisper
from transformers import pipeline

# Load Whisper model
whisper_model = whisper.load_model("base", device="cpu")

# Load the text correction model
correction_pipeline = pipeline("text2text-generation", model="tiiuae/falcon-7b-instruct", device=-1)

# Function to preprocess audio and transcribe it using Whisper
def transcribe_audio(audio_file):
    transcription = whisper_model.transcribe(audio_file)
    return transcription["text"]

# Function to correct grammar in text
def correct_text(raw_text):
    corrected = correction_pipeline(raw_text, max_length=200, num_return_sequences=1)[0]["generated_text"]
    return corrected

# Function to process the pipeline
def process_pipeline(audio_file):
    raw_transcription = transcribe_audio(audio_file)
    corrected_transcription = correct_text(raw_transcription)
    return raw_transcription, corrected_transcription

# Gradio Interface
interface = gr.Interface(
    fn=process_pipeline,
    inputs=gr.Audio(type="filepath", label="Upload Audio"),
    outputs=[
        gr.Textbox(label="Raw Transcription"),
        gr.Textbox(label="Corrected Transcription"),
    ],
    title="Speech Correction Demo",
    description="Upload an audio file to see raw transcription and grammar-corrected output.",
)

# Launch the app
interface.launch()