import gradio as gr
from transformers import pipeline
import pytube

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

def generate(text):
    yt = pytube.YouTube(text)
    title = f"Title: {yt.title.lower()}\nComment:"
    result = generator(title, 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()