Upload 2 files

app.py

@@ -1,43 +1,10 @@
 import streamlit as st
 import torch
-from urllib.parse import urlparse, parse_qs
+#from urllib.parse import urlparse, parse_qs
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-# https://pypi.org/project/youtube-transcript-api/
-from youtube_transcript_api import YouTubeTranscriptApi
 
-def get_video_id(url: str) -> str:
-    """
-    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
-    """
-    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]
-    return None
-
-
-def get_youtube_subtitle(video_id: str) -> str:
-    try:  
-        parse = YouTubeTranscriptApi.get_transcript(video_id, languages=['ru'])
-        result = ''
-        for i in parse:
-            if (i['text'][0] =='[') & (i['text'][-1] ==']'): continue
-            result += ' ' + i['text']
-        result = result.strip()[0].upper() + result.strip()[1:]
-        return result.strip()
-    except:
-        return None
+#from youtube_transcript_api import YouTubeTranscriptApi
+from util import get_video_id, get_youtube_subtitle
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 m_name = '/home/user/app/model'
@@ -48,13 +15,12 @@ model.to(device)
         
 if __name__ == "__main__":
     st.header("Annotation of subtitles from YouTube")
-
-#    url = st.text_input('Enter the URL of the Youtube video', 'https://www.youtube.com/watch?v=HGSVsK32rKA ')
     option = st.selectbox(
-           'Video for example:',
-               ('https://www.youtube.com/watch?v=HGSVsK32rKA',
-               'https://www.youtube.com/watch?v=fSpARfZ3I50',
-               'https://www.youtube.com/watch?v=3lEMopaRSjw'))
+                          'Video for example:',
+                          ('https://www.youtube.com/watch?v=HGSVsK32rKA',
+                           'https://www.youtube.com/watch?v=fSpARfZ3I50',
+                           'https://www.youtube.com/watch?v=3lEMopaRSjw')
+                          )
     url = st.text_input(':green[Enter your URL of the Youtube video] 👇', option)
 
     video_id = get_video_id(url)
@@ -72,7 +38,6 @@ if __name__ == "__main__":
                                 return_tensors="pt",
                                 )["input_ids"]
 
-
             if st.button('Compute summary', help='Click me'):
                 outputs = model.generate(inputs.to(device), max_new_tokens=100, do_sample=False)
                 summary = tokenizer.decode(outputs[0], skip_special_tokens=True)

util.py

@@ -0,0 +1,36 @@
+from urllib.parse import urlparse, parse_qs
+from youtube_transcript_api import YouTubeTranscriptApi
+
+def get_video_id(url: str) -> str:
+    """
+    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
+    """
+    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]
+    return None
+
+
+def get_youtube_subtitle(video_id: str) -> str:
+    try:  
+        parse = YouTubeTranscriptApi.get_transcript(video_id, languages=['ru'])
+        result = ''
+        for i in parse:
+            if (i['text'][0] =='[') & (i['text'][-1] ==']'): continue
+            result += ' ' + i['text']
+        result = result.strip()[0].upper() + result.strip()[1:]
+        return result.strip()
+    except:
+        return None