File size: 1,978 Bytes
df8e313
 
43f8454
9554b3a
 
43f8454
 
9554b3a
64491cc
10933cd
0750508
64491cc
 
 
9554b3a
 
 
7b3ea81
43f8454
 
 
 
 
441212c
ddbbd0b
9554b3a
 
 
 
 
 
b77a531
64491cc
 
 
 
 
 
 
9554b3a
6dd3d81
fe1cfcd
 
 
 
9554b3a
6dd3d81
9554b3a
6dd3d81
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import streamlit as st
import torch
#from urllib.parse import urlparse, parse_qs
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

#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'
#m_name = '../model'
tokenizer = AutoTokenizer.from_pretrained(m_name)
model = AutoModelForSeq2SeqLM.from_pretrained(m_name)
model.to(device)
        
if __name__ == "__main__":
    st.header("Annotation of subtitles from YouTube")
    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')
                          )
    url = st.text_input(':green[Enter your URL of the Youtube video] 👇', option)

    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.markdown(subtitle)

            inputs = tokenizer( [subtitle],
                                max_length=1000,
                                padding="max_length",
                                truncation=True,
                                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)
                st.subheader('Summary')
                st.markdown(summary)
        else:
            st.markdown(':red[Subtitles are disabled for this video]')
    else:
        st.markdown(':red[Video clip is not detected]')