DrishtiSharma commited on
Commit
55dac6c
1 Parent(s): 45c2651

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +120 -0
app.py ADDED
@@ -0,0 +1,120 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from collections import deque
2
+ import streamlit as st
3
+ import torch
4
+ from streamlit_player import st_player
5
+ from transformers import AutoModelForCTC, Wav2Vec2Processor
6
+ from streaming import ffmpeg_stream
7
+
8
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
9
+ player_options = {
10
+ "events": ["onProgress"],
11
+ "progress_interval": 200,
12
+ "volume": 1.0,
13
+ "playing": True,
14
+ "loop": False,
15
+ "controls": False,
16
+ "muted": False,
17
+ "config": {"youtube": {"playerVars": {"start": 1}}},
18
+ }
19
+
20
+ # disable rapid fading in and out on `st.code` updates
21
+ st.markdown("<style>.element-container{opacity:1 !important}</style>", unsafe_allow_html=True)
22
+
23
+ @st.cache(hash_funcs={torch.nn.parameter.Parameter: lambda _: None})
24
+ def load_model(model_path="facebook/wav2vec2-large-robust-ft-swbd-300h"):
25
+ processor = Wav2Vec2Processor.from_pretrained(model_path)
26
+ model = AutoModelForCTC.from_pretrained(model_path).to(device)
27
+ return processor, model
28
+
29
+ model_path = st.radio("Select a model", ("jonatasgrosman/wav2vec2-large-xlsr-53-spanish", "facebook/wav2vec2-large-xlsr-53-spanish", "patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm", "jonatasgrosman/wav2vec2-xls-r-1b-spanish", "glob-asr/xls-r-es-test-lm"))
30
+
31
+ processor, model = load_model(model_path)
32
+
33
+ def stream_text(url, chunk_duration_ms, pad_duration_ms):
34
+ sampling_rate = processor.feature_extractor.sampling_rate
35
+
36
+ # calculate the length of logits to cut from the sides of the output to account for input padding
37
+ output_pad_len = model._get_feat_extract_output_lengths(int(sampling_rate * pad_duration_ms / 1000))
38
+
39
+ # define the audio chunk generator
40
+ stream = ffmpeg_stream(url, sampling_rate, chunk_duration_ms=chunk_duration_ms, pad_duration_ms=pad_duration_ms)
41
+
42
+ leftover_text = ""
43
+ for i, chunk in enumerate(stream):
44
+ input_values = processor(chunk, sampling_rate=sampling_rate, return_tensors="pt").input_values
45
+
46
+ with torch.no_grad():
47
+ logits = model(input_values.to(device)).logits[0]
48
+ if i > 0:
49
+ logits = logits[output_pad_len : len(logits) - output_pad_len]
50
+ else: # don't count padding at the start of the clip
51
+ logits = logits[: len(logits) - output_pad_len]
52
+
53
+ predicted_ids = torch.argmax(logits, dim=-1).cpu().tolist()
54
+ if processor.decode(predicted_ids).strip():
55
+ leftover_ids = processor.tokenizer.encode(leftover_text)
56
+ # concat the last word (or its part) from the last frame with the current text
57
+ text = processor.decode(leftover_ids + predicted_ids)
58
+ # don't return the last word in case it's just partially recognized
59
+ text, leftover_text = text.rsplit(" ", 1)
60
+ yield text
61
+ else:
62
+ yield leftover_text
63
+ leftover_text = ""
64
+ yield leftover_text
65
+
66
+ def main():
67
+ state = st.session_state
68
+ st.header("Video ASR Streamlit from Youtube Link")
69
+
70
+ with st.form(key="inputs_form"):
71
+
72
+ # Our worlds best teachers on subjects of AI, Cognitive, Neuroscience for our Behavioral and Medical Health
73
+ ytSeanKelly="https://www.youtube.com/watch?v=85n-EEqKMSY"
74
+ ytSamHarris="https://www.youtube.com/watch?v=4dC_nRYIDZU&list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&index=2"
75
+ ytJohnAbramson="https://www.youtube.com/watch?v=arrokG3wCdE&list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&index=3"
76
+ ytElonMusk="https://www.youtube.com/watch?v=DxREm3s1scA&list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&index=4"
77
+ ytJeffreyShainline="https://www.youtube.com/watch?v=EwueqdgIvq4&list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&index=5"
78
+ ytJeffHawkins="https://www.youtube.com/watch?v=Z1KwkpTUbkg&list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&index=6"
79
+ ytSamHarris="https://youtu.be/Ui38ZzTymDY?list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L"
80
+ ytSamHarris="https://youtu.be/4dC_nRYIDZU?list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&t=7809"
81
+ ytSamHarris="https://youtu.be/4dC_nRYIDZU?list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&t=7809"
82
+ ytSamHarris="https://youtu.be/4dC_nRYIDZU?list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&t=7809"
83
+ state.youtube_url = st.text_input("YouTube URL", ytSeanKelly)
84
+
85
+
86
+ state.chunk_duration_ms = st.slider("Audio chunk duration (ms)", 2000, 10000, 3000, 100)
87
+ state.pad_duration_ms = st.slider("Padding duration (ms)", 100, 5000, 1000, 100)
88
+ submit_button = st.form_submit_button(label="Submit")
89
+
90
+ if submit_button or "asr_stream" not in state:
91
+ # a hack to update the video player on value changes
92
+ state.youtube_url = (
93
+ state.youtube_url.split("&hash=")[0]
94
+ + f"&hash={state.chunk_duration_ms}-{state.pad_duration_ms}"
95
+ )
96
+ state.asr_stream = stream_text(
97
+ state.youtube_url, state.chunk_duration_ms, state.pad_duration_ms
98
+ )
99
+ state.chunks_taken = 0
100
+
101
+
102
+ state.lines = deque([], maxlen=100) # limit to the last n lines of subs
103
+
104
+
105
+ player = st_player(state.youtube_url, **player_options, key="youtube_player")
106
+
107
+ if "asr_stream" in state and player.data and player.data["played"] < 1.0:
108
+ # check how many seconds were played, and if more than processed - write the next text chunk
109
+ processed_seconds = state.chunks_taken * (state.chunk_duration_ms / 1000)
110
+ if processed_seconds < player.data["playedSeconds"]:
111
+ text = next(state.asr_stream)
112
+ state.lines.append(text)
113
+ state.chunks_taken += 1
114
+ if "lines" in state:
115
+ # print the lines of subs
116
+ st.code("\n".join(state.lines))
117
+
118
+
119
+ if __name__ == "__main__":
120
+ main()