import gradio as gr
import openai
import yt_dlp
import os
openai.api_key = os.environ['OPENAI_API_KEY']

def asr(url):
    # download audio
    # Options for youtube-dl
    ydl_opts = {
        'outtmpl': 'my_video.mp4',
        'overwrites': True, 
    }

    # Create a youtube-dl object
    ydl = yt_dlp.YoutubeDL(ydl_opts)

    # Download the video
    info_dict = ydl.extract_info(url, download=True)

    audio_file= open("my_video.mp4", "rb")
    transcript = openai.Audio.transcribe("whisper-1", audio_file)
    output = openai.ChatCompletion.create(
    model="gpt-3.5-turbo",
    messages=[
            {"role": "user", "content": "Transcript: {transcript}. \n Translate the video conversation transcript into fluent Chinese. Chinese: ".format(transcript=transcript["text"])},
        ]
    )
    return output['choices'][0]['message']['content'], transcript["text"]

title = """
韬译云间"""
# Create an instruction input component
instruction = """
<div style="border: 2px solid #000; padding: 10px; border-radius: 5px;">
一键输入视频链接，轻松实现中文翻译，畅享视频无障碍沟通 <span style="color: grey;">-- powered by OpenAI Whisper & ChatGPT.</span>.<br>

1.将视频链接（支持Twitter、YouTube）复制粘贴至输入框，点击提交（Submit）即可;
2.为保证翻译质量，目前仅支持处理时长不超过5分钟的短视频。
</div>"""
# Create a text input component
text_input = gr.inputs.Textbox()

demo = gr.Interface(fn=asr, 
                    inputs=gr.inputs.Textbox(label="粘贴视频链接"),
                    outputs=[
                        gr.outputs.Textbox(label="中文"),
                        gr.outputs.Textbox(label="英文"),
                    ],
                    title=title,
                    description=instruction)

demo.launch()