from transformers import pipeline
from pytube import YouTube

import gradio as gr
import librosa

import hopsworks

project = hopsworks.login()
fs = project.get_feature_store()

dataset_api = project.get_dataset_api()

dataset_api.download("Resources/titanic/images/latest_titanic.png", overwrite=True)  # change link

# pipe = pipeline(model="fimster/whisper-small-sv-SE") # change model
# pipe = pipeline(model="ayberkuckun/whisper-small-sv-SE")
pipe = pipeline(model="openai/whisper-small", task="automatic-speech-recognition", chunk_length_s=30)


def transcribe(url):
    selected_video = YouTube(url)

    try:
        audio = selected_video.streams.filter(only_audio=True)[0]
    except:
        raise Exception("Can't find an mp4 audio.")

    audio.download(filename="audio.mp3")

    speech_array, _ = librosa.load("audio.mp3", sr=16000)

    output = pipe(speech_array)

    return output["text"]


iface = gr.Interface(
    fn=transcribe,
    inputs=gr.Textbox("https://www.youtube.com/watch?v=n9g12Xm9UJM", label="Paste a YouTube video URL"),
    outputs=[gr.Textbox(label="Transcription")],
    title="Whisper Small Swedish",
    description="Realtime demo for Swedish speech recognition using a fine-tuned Whisper small model.",
    allow_flagging="never"
)

iface.launch()
#              gr.Image("latest_titanic.png", label="Model Scores")