Upload app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st 
 from urllib.parse import urlparse, parse_qs
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 # https://pypi.org/project/youtube-transcript-api/
 from youtube_transcript_api import YouTubeTranscriptApi
 
@@ -23,7 +23,6 @@ def get_video_id(url: str) -> str:
             return query.path.split('/')[2]
         if query.path[:3] == '/v/':
             return query.path.split('/')[2]
-    # fail?
     return None
 
 
@@ -39,30 +38,44 @@ def get_youtube_subtitle(video_id: str) -> str:
     except:
         return None
 
-st.header("Annotation of subtitles from YouTube")
-st.text('Load model...')
-m_name = 'summarize1'
-tokenizer = AutoTokenizer.from_pretrained(m_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(m_name)
-st.text('Model is loaded')
+        
+if __name__ == "__main__":
+    st.header("Annotation of subtitles from YouTube")
+    # st.text('Load model...')
+    # m_name = '/content/drive/MyDrive/Colab Notebooks/Netology/diplom_neto/summarize1'
+    m_name = "csebuetnlp/mT5_multilingual_XLSum"
+    # 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=HGSVsK32rKA')
-video_id = get_video_id(url)
+    url = st.text_input('Enter the URL of the Youtube video', 'https://www.youtube.com/watch?v=HGSVsK32rKA')
+    video_id = get_video_id(url)
 
-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...')
+    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...')
 
-        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)
+            # inputs = tokenizer(
+            #                     [subtitle],
+            #                     max_length=600,
+            #                     padding="max_length",
+            #                     truncation=True,
+            #                     return_tensors="pt",
+            #                     )["input_ids"]
+
+            # # inputs = tokenizer(subtitle, 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)
+            translator = pipeline("summarization", model=m_name, 
+                                  tokenizer=m_name, max_length=100, device=0
+                                  )
+
+            st.subheader('Summary')
+            st.text(translator(subtitle))
+        else:
+            st.write('Subtitles are disabled for this video')
     else:
-        st.write('Subtitles are disabled for this video')
-else:
-    st.write('Video clip is not detected')
+        st.write('Video clip is not detected')