added ui and refactoring

.gitignore

@@ -0,0 +1 @@
+.idea/

.streamlit/config.toml

@@ -0,0 +1,4 @@
+[theme]
+backgroundColor="#d2e7ee"
+secondaryBackgroundColor="#79c1ee"
+textColor="#151516"

app.py

@@ -1,22 +1,20 @@
 import os
 
-
 import streamlit as st
 from urllib.parse import urlparse, parse_qs
 
-from tqdm import tqdm
 from stqdm import stqdm
 
 # https://github.com/pytorch/pytorch/issues/77764
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
 
-from youtube_transcript_api import YouTubeTranscriptApi
+from youtube_transcript_api import YouTubeTranscriptApi, NoTranscriptFound, TranscriptsDisabled
 
-from transformers import pipeline, T5ForConditionalGeneration, T5Tokenizer
+from transformers import pipeline
 
 import torch
 
-# Setting device for PYTorch
+# Setting device for PyTorch
 if torch.cuda.is_available():
     device = torch.device('cuda')
 elif torch.has_mps:
@@ -25,47 +23,79 @@ else:
     device = torch.device('cpu')
 
 
+class InvalidURLException(Exception):
+    pass
+
 
-def get_videoid_from_url(url:str):
+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)
 
-    try:
+    if ('v' in query) & ('youtube.com' in url_data.netloc):
         video_id = query["v"][0]
-    except KeyError:
-        video_id = ''
+    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 process_click_callback():
-    st.session_state.process_btn = True
 
-    print('Using {} device'.format(device))
+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('aircAruvnKk')  # 3blue1Brown
+    transcript_list = YouTubeTranscriptApi.list_transcripts(video_id)
 
     try:
         transcript = transcript_list.find_manually_created_transcript(['en'])
-    except Exception as e:
-        print('No manual transcripts were found, trying to load generated ones...')
+    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]']
-    subtitles_len = [len(sbt) for sbt in subtitles]
-    sbt_mean_len = sum(subtitles_len)/len(subtitles_len)
 
-    print('Mean length of subtitles: {}'.format(sbt_mean_len))
-    print(subtitles)
-    print(len(subtitles))
+    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))
-    print('Number subtitles per summary: {}'.format(n_sbt_per_step))
 
     n_steps = len(subtitles) // n_sbt_per_step if len(subtitles) % n_sbt_per_step == 0 else \
         len(subtitles) // n_sbt_per_step + 1
@@ -73,9 +103,7 @@ def process_click_callback():
     summaries = []
 
     for i in stqdm(range(n_steps)):
-        sbt_txt = ' '.join(subtitles[n_sbt_per_step*i:n_sbt_per_step*(i+1)])
-        # print('length of text: {}'.format(len(sbt_txt)))
-        # print(sbt_txt)
+        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)
@@ -83,44 +111,59 @@ def process_click_callback():
         summary = summarizer(sbt_txt, do_sample=False)
         summary = summary[0]['summary_text']
 
-        # print('Summary: ' + summary)
         summaries.append(summary)
 
-    out = ' '.join(summaries)
-    print(out)
+    return ' '.join(summaries)
 
-    st.session_state.summary_output = out
+
+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="✅")
-    st.session_state.process_btn = False
 
+    is_processing = False
 
 
-def main():
+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=aircAruvnKk')
-    st.code('https://youtu.be/p0G68ORc8uQ')
+    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 = st.text_input('YouTube Video URL:',  placeholder='YouTube URL',
-                                 label_visibility='collapsed')
-        st.write(get_videoid_from_url(video_url))
+        video_url_inp = st.text_input('YouTube Video URL:', placeholder='YouTube URL',
+                                      label_visibility='collapsed')
 
     with col2:
-        st.button('Process 📭', key='process_btn', on_click=process_click_callback)
-
-    st.text_area(label='', key='summary_output', height=444)
-
+        process_btn = st.button('🗜️Process', key='process_btn', on_click=process_click_callback)
 
-
-
-
-
-    # x = st.slider('Select a value')
-    # st.write(x, 'squared is', x * x)
-
-
-if __name__ == "__main__":
-    main()
+    summary_out_txt = st.text_area(label='', key='summary_output', height=400)

requirements.txt

@@ -2,4 +2,5 @@ torch
 transformers
 youtube_transcript_api
 tqdm
-stqdm
+stqdm
+streamlit