from transformers import pipeline
import gradio as gr
from pytube import YouTube
import os

pipe = pipeline(model="torileatherman/train_first_try")  # change to "your-username/the-name-you-picked"

def transcribe(audio):
    text = pipe(audio)["text"]
    return text

def transcribe_url(url):
    youtube = YouTube(str(url))
    audio = youtube.streams.filter(only_audio=True).first().download('yt_video')
    text = pipe(audio)["text"]
    return text

url_demo = gr.Interface(
    fn=transcribe_url, 
    inputs="text", 
    outputs="text",
    title="Whisper Swedish",
    description="Swedish speech and audio recognition using a fine-tuned Whisper small model",
)

voice_demo = gr.Interface(
    fn=transcribe, 
    inputs=gr.Audio(source="microphone", type="filepath"), 
    outputs="text",
    title="Whisper Swedish",
    description="Swedish speech and audio recognition using a fine-tuned Whisper small model",
)

demo = gr.TabbedInterface([url_demo, voice_demo], ["YouTube Video to Text", "Audio to Text"])

demo.launch()