import gradio as gr
import pytube
from transformers import pipeline
from time import sleep


generator = pipeline('text-generation', model='parseny/youtube_comment_generation_02')

def generate(text):
    yt = pytube.YouTube(text)
    while True:
        try:
            title = yt.title
            break
        except:
            sleep(1)
            yt = pytube.YouTube(text)
            continue
    title = title.lower()
    title = title.replace("[^a-zA-Z#]", " ")
    title = title.replace(" +", " ")
    prompt = f"Title: {title}\nComment:"
    result = generator(prompt, max_new_tokens=30, num_return_sequences=1, do_sample=True)
    return result[0]["generated_text"]

examples = [
    ["https://www.youtube.com/watch?v=mCV44C5rQ2M"],
    ["https://www.youtube.com/watch?v=sitXeGjm4Mc"],
    ["https://www.youtube.com/watch?v=oQS8KUoWL8E"],
]

demo = gr.Interface(
    fn=generate,
    inputs=gr.inputs.Textbox(lines=5, label="Input Text"),
    outputs=gr.outputs.Textbox(label="Generated Text"),
    examples=examples
)

demo.launch()