import whisper
import gradio as gr


from transformers import pipeline

# Force the model to run on CPU
device = "cpu"
print("Running on CPU")


# Load the tiny Whisper model
whisper_model = whisper.load_model("tiny")

# Load the text summarization model from Hugging Face
#summarizer = pipeline(task="summarization", model="facebook/bart-large-cnn")
#summarizer = pipeline("summarization", model="allenai/led-base-16384")
summarizer = pipeline(task="summarization", model="facebook/bart-large-cnn")


# Function to transcribe and summarize the audio file
def transcribe_and_summarize(audio):
    # Step 1: Transcribe the audio using Whisper
    transcription_result = whisper_model.transcribe(audio)
    transcription = transcription_result['text']

    # Step 2: Summarize the transcription
    summary = summarizer(transcription, min_length=50, max_length=200)

    summary_text = summary[0]['summary_text']

    return transcription, summary_text

# Create the Gradio interface
demo = gr.Interface(
    fn=transcribe_and_summarize,  # The function to be called for transcription
    inputs=gr.Audio(type="filepath", label="Upload your audio file"),  # Input audio field
    outputs=[gr.Textbox(label="Transcription"), gr.Textbox(label="Summary")],  # Output fields
    title="Whisper Tiny Transcription and Summarization",
    examples=["Classification_and_Regression_in_Machine_Learning.mp3"],
    description="Upload an audio file, get the transcription from Whisper tiny model and the summarized version using Hugging Face."
)

# Launch the Gradio interface
demo.launch(debug=True)