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


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

def generate(text):
    yt = pytube.YouTube(text)
    while True:
        try:
            title = yt.title
            title = title.replace("[^a-zA-Z#]", " ")
            title = title.replace(" +", " ")
            title = title.lower()
            prompt = f"Title: {title}\nComment:"
            break
        except:
            sleep(1)
            yt = pytube.YouTube(text)
            continue
    result = generator(prompt, max_new_tokens=30, pad_token_id=50256, 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"],
]

# input_textbox = gr.inputs.Textbox(lines=5, label="Input Text")
# output_textbox = gr.outputs.Textbox(label="Generated Text")

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

# 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()