effluxriad/YouTube-comments-generator

1import streamlit as st2from PIL import Image3import yt_dlp4from transformers import pipeline5 6# load model7comments_generator = pipeline(task='text-generation', model='gpt2')8 9 10def get_yt_video_title(yt_video_link: str) -> str:11    ydl = yt_dlp.YoutubeDL({'outtmpl': '%(id)s.%(ext)s'})12    with ydl:13        try:14            info_dict = ydl.extract_info(yt_video_link, download=False)15            video_title = info_dict['title']16            return video_title17        except yt_dlp.utils.DownloadError as err:18            raise err19 20 21def invalid_yt_link_error():22    st.write("Entered link is incorrect. Please, enter a valid YouTube video link")23 24 25def generate_comment(yt_video_link: str, comments_cnt: int = 1, num_beams: int = 2, length_penalty: int = 5):26    try:27        video_title = get_yt_video_title(yt_video_link)28    except yt_dlp.utils.DownloadError:29        invalid_yt_link_error()30        return31 32    req_text = f"The YouTube video named '{video_title}' may have the following comment: "33    generator_output = comments_generator(req_text,34                                          do_sample=True,35                                          num_return_sequences=comments_cnt,36                                          num_beams=num_beams,37                                          length_penalty=length_penalty)38 39    st.markdown("### Generated comments")40 41    comments_str = ""42    for i in range(comments_cnt):43        comment_str = generator_output[i]['generated_text'].replace(req_text, '').replace('"', '')44        comments_str += str(i + 1) + ". " + comment_str + "\n\n"45 46    st.write(comments_str)47 48 49# ----- application ------50st.title("YouTube comments generator")51 52image = Image.open('img/youtube-comments-img.jpeg')53st.image(image, width=200)54 55st.markdown("### Generation parameters")56with st.form("Comment generating"):57    user_video_link = st.text_input("Insert YouTube video link here")58 59    comments_cnt_col, num_beams_col, len_penalty_col = st.columns(3)60 61    with comments_cnt_col:62        user_comments_cnt = st.slider("Comments to generate:", 1, 10, 1)63 64    with num_beams_col:65        user_num_beams = st.slider("Number of beams:", 1, 10, 1)66 67    with len_penalty_col:68        user_len_penalty = st.slider("Penalty length:", 0, 10, 1)69 70    generate_button_res = st.form_submit_button("Generate comments")71 72if generate_button_res:73    generate_comment(user_video_link, user_comments_cnt, user_num_beams, user_len_penalty)74