import os import streamlit as st from urllib.parse import urlparse, parse_qs from stqdm import stqdm # https://github.com/pytorch/pytorch/issues/77764 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" from youtube_transcript_api import YouTubeTranscriptApi, NoTranscriptFound, TranscriptsDisabled from transformers import pipeline import torch # Setting device for PyTorch if torch.cuda.is_available(): device = torch.device('cuda') elif torch.has_mps: device = torch.device('mps') else: device = torch.device('cpu') class InvalidURLException(Exception): pass def get_videoid_from_url(url: str): ''' Gets video ID from give YouTube video URL :param url: YouTube video URL in 2 formats (standard and short) :return: id of YouTube video :raises InvalidURLException: If URL is not valid ''' url_data = urlparse(url) query = parse_qs(url_data.query) if ('v' in query) & ('youtube.com' in url_data.netloc): video_id = query["v"][0] elif 'youtu.be' in url_data.netloc: path_lst = url.split('/') if path_lst: video_id = path_lst[-1] else: raise InvalidURLException('Invalid URL') else: raise InvalidURLException('Invalid URL') return video_id def get_transcripts(url: str): ''' Loads transcripts for given URL :param url: YouTube video URL :return: list, list of subtitles ''' video_id = get_videoid_from_url(video_url_inp) transcript_list = YouTubeTranscriptApi.list_transcripts(video_id) try: transcript = transcript_list.find_manually_created_transcript(['en']) except NoTranscriptFound as e: st.info('No manual transcripts were found, trying to load generated ones...') transcript = transcript_list.find_generated_transcript(['en']) subtitles = transcript.fetch() subtitles = [sbt['text'] for sbt in subtitles if sbt['text'] != '[Music]'] return subtitles def generate_summary(subtitles: list): ''' Creates summary based on subtitles of YouTube video. Uses T5-small model which shows best results for different topics of videos. :param subtitles: list of subtitles strings :return: summary based on subtitles ''' subtitles_len = [len(sbt) for sbt in subtitles] sbt_mean_len = sum(subtitles_len) / len(subtitles_len) # Number of subtitles per step/summary # Since number length of transcripts differs # between generated and manual ones # we set different step size n_sbt_per_step = int(400 / (sbt_mean_len / 4)) n_steps = len(subtitles) // n_sbt_per_step if len(subtitles) % n_sbt_per_step == 0 else \ len(subtitles) // n_sbt_per_step + 1 summaries = [] for i in stqdm(range(n_steps)): sbt_txt = ' '.join(subtitles[n_sbt_per_step * i:n_sbt_per_step * (i + 1)]) summarizer = pipeline('summarization', model='t5-small', tokenizer='t5-small', max_length=512, truncation=True) summary = summarizer(sbt_txt, do_sample=False) summary = summary[0]['summary_text'] summaries.append(summary) return ' '.join(summaries) def process_click_callback(): ''' Callback for process button click ''' global is_processing if is_processing: return else: is_processing = True global video_url_inp try: subtitles = get_transcripts(video_url_inp) except InvalidURLException as iue: is_processing = False st.error('Invalid YouTube URL, please provide URL in format that is shown on Examples') st.experimental_rerun() except TranscriptsDisabled as tde: is_processing = False st.error('Could not retrieve a transcript for given ID') st.experimental_rerun() summary = generate_summary(subtitles) st.session_state.summary_output = summary st.success('Processing complete!', icon="✅") is_processing = False if __name__ == "__main__": # State of processing is_processing = False st.title('YouTube Video Summary 📃') st.markdown('Creates summary for given YouTube video URL based on transcripts.') st.code('https://www.youtube.com/watch?v=skl4OXNA12U') st.code('https://youtu.be/mEQc-iAbEBk') col1, col2 = st.columns(2) with col1: video_url_inp = st.text_input('YouTube Video URL:', placeholder='YouTube URL', label_visibility='collapsed') with col2: process_btn = st.button('🗜️Process', key='process_btn', on_click=process_click_callback) summary_out_txt = st.text_area(label='', key='summary_output', height=400)