Upload app.py

app.py

@@ -1,47 +1,68 @@
-import streamlit as st
+import streamlit as st 
+from urllib.parse import urlparse, parse_qs
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 # https://pypi.org/project/youtube-transcript-api/
 from youtube_transcript_api import YouTubeTranscriptApi
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-
 
 def get_video_id(url: str) -> str:
     """
-    https://github.com/KennethSC/YouTube-Captions-Formatter/blob/master/YouTube_Captions.py
+    Examples:
+    - http://youtu.be/SA2iWivDJiE
+    - http://www.youtube.com/watch?v=_oPAwA_Udwc&feature=feedu
+    - http://www.youtube.com/embed/SA2iWivDJiE
+    - http://www.youtube.com/v/SA2iWivDJiE?version=3&hl=en_US
     """
-    split_at = 'watch?v='
-    if 'https://www.youtube.com/watch?v=' in url:
-        video_id = url.partition(split_at)[2]
-    else:
-        raise Exception("This is not a valid video URL")
-    return video_id
+    query = urlparse(url)
+    if query.hostname == 'youtu.be':
+        return query.path[1:]
+    if query.hostname in ('www.youtube.com', 'youtube.com'):
+        if query.path == '/watch':
+            p = parse_qs(query.query)
+            return p['v'][0]
+        if query.path[:7] == '/embed/':
+            return query.path.split('/')[2]
+        if query.path[:3] == '/v/':
+            return query.path.split('/')[2]
+    # fail?
+    return None
 
 
 def get_youtube_subtitle(video_id: str) -> str:
-    try:
+    try:  
         parse = YouTubeTranscriptApi.get_transcript(video_id, languages=['ru'])
         result = ''
         for i in parse:
-            if (i['text'][0] =='[') & (i['text'][-1] ==']'):
-                continue
+            if (i['text'][0] =='[') & (i['text'][-1] ==']'): continue
             result += ' ' + i['text']
-        return result
+        result = result.strip()[0].upper() + result.strip()[1:]
+        return result.strip()
     except:
-        raise Exception("This video doesn't have a captions transcript")
+        return None
 
+st.header("Annotation of subtitles from YouTube")
+st.text('Load model...')
+m_name = '/content/drive/MyDrive/Colab Notebooks/Netology/diplom_neto/summarize1'
+tokenizer = AutoTokenizer.from_pretrained(m_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(m_name)
+st.text('Model is loaded')
 
-url = st.text_input('Enter the URL of the Youtube video', 'https://www.youtube.com/watch?v=yR4VmxwZh0s')
+url = st.text_input('Enter the URL of the Youtube video', 'https://www.youtube.com/watch?v=HGSVsK32rKA')
 video_id = get_video_id(url)
-subtitle = get_youtube_subtitle(video_id)
-st.write('Video_id', video_id)
 
-st.text(subtitle)
-m_name = 'summarize1'
-tokenizer = AutoTokenizer.from_pretrained(m_name)
-inputs = tokenizer(subtitle[:1024], return_tensors="pt").input_ids
+if video_id is not None:
+    subtitle = get_youtube_subtitle(video_id)
+    if subtitle is not None:
+        st.subheader('Subtitles')
+        st.text(subtitle)
+        st.text('Compute summary...')
 
-model = AutoModelForSeq2SeqLM.from_pretrained(m_name)
-outputs = model.generate(inputs, max_new_tokens=100, do_sample=False)
-
-summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
-st.write('subtitle', "")
-st.text(summary)
+        inputs = tokenizer(subtitle[:1024], return_tensors="pt").input_ids
+        outputs = model.generate(inputs, max_new_tokens=100, do_sample=False)
+        summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        st.subheader('Summary')
+        st.text(summary)
+    else:
+        st.write('Subtitles are disabled for this video')
+else:
+    st.write('Video clip is not detected')