Spaces:
Runtime error
Runtime error
Commit
•
ceb338c
1
Parent(s):
f18ba93
Create new file
Browse files
app.py
ADDED
@@ -0,0 +1,119 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
+
processor, model = load_model()
|
30 |
+
|
31 |
+
def stream_text(url, chunk_duration_ms, pad_duration_ms):
|
32 |
+
sampling_rate = processor.feature_extractor.sampling_rate
|
33 |
+
|
34 |
+
# calculate the length of logits to cut from the sides of the output to account for input padding
|
35 |
+
output_pad_len = model._get_feat_extract_output_lengths(int(sampling_rate * pad_duration_ms / 1000))
|
36 |
+
|
37 |
+
# define the audio chunk generator
|
38 |
+
stream = ffmpeg_stream(url, sampling_rate, chunk_duration_ms=chunk_duration_ms, pad_duration_ms=pad_duration_ms)
|
39 |
+
|
40 |
+
leftover_text = ""
|
41 |
+
for i, chunk in enumerate(stream):
|
42 |
+
input_values = processor(chunk, sampling_rate=sampling_rate, return_tensors="pt").input_values
|
43 |
+
|
44 |
+
with torch.no_grad():
|
45 |
+
logits = model(input_values.to(device)).logits[0]
|
46 |
+
if i > 0:
|
47 |
+
logits = logits[output_pad_len : len(logits) - output_pad_len]
|
48 |
+
else: # don't count padding at the start of the clip
|
49 |
+
logits = logits[: len(logits) - output_pad_len]
|
50 |
+
|
51 |
+
predicted_ids = torch.argmax(logits, dim=-1).cpu().tolist()
|
52 |
+
if processor.decode(predicted_ids).strip():
|
53 |
+
leftover_ids = processor.tokenizer.encode(leftover_text)
|
54 |
+
# concat the last word (or its part) from the last frame with the current text
|
55 |
+
text = processor.decode(leftover_ids + predicted_ids)
|
56 |
+
# don't return the last word in case it's just partially recognized
|
57 |
+
text, leftover_text = text.rsplit(" ", 1)
|
58 |
+
yield text
|
59 |
+
else:
|
60 |
+
yield leftover_text
|
61 |
+
leftover_text = ""
|
62 |
+
yield leftover_text
|
63 |
+
|
64 |
+
def main():
|
65 |
+
state = st.session_state
|
66 |
+
st.header("Video ASR Streamlit from Youtube Link")
|
67 |
+
|
68 |
+
with st.form(key="inputs_form"):
|
69 |
+
|
70 |
+
# Our worlds best teachers on subjects of AI, Cognitive, Neuroscience for our Behavioral and Medical Health
|
71 |
+
ytJoschaBach="https://youtu.be/cC1HszE5Hcw?list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&t=8984"
|
72 |
+
ytSamHarris="https://www.youtube.com/watch?v=4dC_nRYIDZU&list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&index=2"
|
73 |
+
ytJohnAbramson="https://www.youtube.com/watch?v=arrokG3wCdE&list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&index=3"
|
74 |
+
ytElonMusk="https://www.youtube.com/watch?v=DxREm3s1scA&list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&index=4"
|
75 |
+
ytJeffreyShainline="https://www.youtube.com/watch?v=EwueqdgIvq4&list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&index=5"
|
76 |
+
ytJeffHawkins="https://www.youtube.com/watch?v=Z1KwkpTUbkg&list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&index=6"
|
77 |
+
ytSamHarris="https://youtu.be/Ui38ZzTymDY?list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L"
|
78 |
+
ytSamHarris="https://youtu.be/4dC_nRYIDZU?list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&t=7809"
|
79 |
+
ytSamHarris="https://youtu.be/4dC_nRYIDZU?list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&t=7809"
|
80 |
+
ytSamHarris="https://youtu.be/4dC_nRYIDZU?list=PLHgX2IExbFouJoqEr8JMF5MbZSbyC91-L&t=7809"
|
81 |
+
ytTimelapseAI="https://www.youtube.com/watch?v=63yr9dlI0cU&list=PLHgX2IExbFovQybyfltywXnqZi5YvaSS-"
|
82 |
+
state.youtube_url = st.text_input("YouTube URL", ytTimelapseAI)
|
83 |
+
|
84 |
+
|
85 |
+
state.chunk_duration_ms = st.slider("Audio chunk duration (ms)", 2000, 10000, 3000, 100)
|
86 |
+
state.pad_duration_ms = st.slider("Padding duration (ms)", 100, 5000, 1000, 100)
|
87 |
+
submit_button = st.form_submit_button(label="Submit")
|
88 |
+
|
89 |
+
if submit_button or "asr_stream" not in state:
|
90 |
+
# a hack to update the video player on value changes
|
91 |
+
state.youtube_url = (
|
92 |
+
state.youtube_url.split("&hash=")[0]
|
93 |
+
+ f"&hash={state.chunk_duration_ms}-{state.pad_duration_ms}"
|
94 |
+
)
|
95 |
+
state.asr_stream = stream_text(
|
96 |
+
state.youtube_url, state.chunk_duration_ms, state.pad_duration_ms
|
97 |
+
)
|
98 |
+
state.chunks_taken = 0
|
99 |
+
|
100 |
+
|
101 |
+
state.lines = deque([], maxlen=100) # limit to the last n lines of subs
|
102 |
+
|
103 |
+
|
104 |
+
player = st_player(state.youtube_url, **player_options, key="youtube_player")
|
105 |
+
|
106 |
+
if "asr_stream" in state and player.data and player.data["played"] < 1.0:
|
107 |
+
# check how many seconds were played, and if more than processed - write the next text chunk
|
108 |
+
processed_seconds = state.chunks_taken * (state.chunk_duration_ms / 1000)
|
109 |
+
if processed_seconds < player.data["playedSeconds"]:
|
110 |
+
text = next(state.asr_stream)
|
111 |
+
state.lines.append(text)
|
112 |
+
state.chunks_taken += 1
|
113 |
+
if "lines" in state:
|
114 |
+
# print the lines of subs
|
115 |
+
st.code("\n".join(state.lines))
|
116 |
+
|
117 |
+
|
118 |
+
if __name__ == "__main__":
|
119 |
+
main()
|