from transformers import pipeline
import gradio as gr
from youtube_dl import YoutubeDL


# Function to download audio from YouTube link
def download_audio(youtube_link, output_path):
    ydl_opts = {
        "format": "bestaudio/best",
        "outtmpl": output_path,
        "postprocessors": [
            {
                "key": "FFmpegExtractAudio",
                "preferredcodec": "mp3",
                "preferredquality": "192",
            }
        ],
    }
    with YoutubeDL(ydl_opts) as ydl:
        ydl.download([youtube_link])


# Function to transcribe audio
def transcribe_audio(audio_path):
    model = pipeline(model="SofiaK/checkpoints")
    return model(audio_path)["text"]


interface = gr.Interface(
    fn=lambda input_type, audio_or_link: transcribe_audio(audio_or_link)
    if input_type == "audio"
    else transcribe_audio(download_audio(audio_or_link, "temp.mp3")),
    inputs=[
        gr.Radio(["audio", "youtube"], label="Select Input Type"),
        gr.Audio(
            sources=["upload", "microphone"],
            type="filepath",
            label="Upload Audio, or speak in the microphone",
        ),
        gr.Textbox(default="https://www.youtube.com/", label="Youtube Link"),
    ],
    outputs=gr.Text(label="Model output"),
    title="Whisper-RU",
    description="Fine-tuned Whisper for Russian language",
)

interface.launch(share=True)