Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import numpy as np | |
| import os, time, librosa, torch | |
| from pyannote.audio import Pipeline | |
| from transformers import pipeline | |
| from utils import second_to_timecode, download_from_youtube | |
| MODEL_NAME = 'bayartsogt/whisper-large-v2-mn-13' | |
| lang = 'mn' | |
| chunk_length_s = 9 | |
| vad_activation_min_duration = 9 # sec | |
| device = 0 if torch.cuda.is_available() else "cpu" | |
| SAMPLE_RATE = 16_000 | |
| ######## LOAD MODELS FROM HUB ######## | |
| dia_model = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=os.environ['TOKEN']) | |
| vad_model = Pipeline.from_pretrained("pyannote/voice-activity-detection", use_auth_token=os.environ['TOKEN']) | |
| pipe = pipeline(task="automatic-speech-recognition", model=MODEL_NAME, chunk_length_s=chunk_length_s, device=device) | |
| pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language=lang, task="transcribe") | |
| print("----------> Loaded models <-----------") | |
| def generator(youtube_link, microphone, file_upload, num_speakers, max_duration, history): | |
| if int(youtube_link != '') + int(microphone is not None) + int(file_upload is not None) != 1: | |
| raise Exception(f"Only one of the source should be given youtube_link={youtube_link}, microphone={microphone}, file_upload={file_upload}") | |
| history = history or "" | |
| if microphone: | |
| path = microphone | |
| elif file_upload: | |
| path = file_upload | |
| elif youtube_link: | |
| path = download_from_youtube(youtube_link) | |
| waveform, sampling_rate = librosa.load(path, sr=SAMPLE_RATE, mono=True, duration=max_duration) | |
| print(waveform.shape, sampling_rate) | |
| waveform_tensor = torch.unsqueeze(torch.tensor(waveform), 0).to(device) | |
| dia_result = dia_model({ | |
| "waveform": waveform_tensor, | |
| "sample_rate": sampling_rate, | |
| }, num_speakers=num_speakers) | |
| counter = 1 | |
| for speech_turn, track, speaker in dia_result.itertracks(yield_label=True): | |
| print(f"{speech_turn.start:4.1f} {speech_turn.end:4.1f} {speaker}") | |
| _start = int(sampling_rate * speech_turn.start) | |
| _end = int(sampling_rate * speech_turn.end) | |
| data = waveform[_start: _end] | |
| if speech_turn.end - speech_turn.start > vad_activation_min_duration: | |
| print(f'audio duration {speech_turn.end - speech_turn.start} sec ----> activating VAD') | |
| vad_output = vad_model({ | |
| 'waveform': waveform_tensor[:, _start:_end], | |
| 'sample_rate': sampling_rate}) | |
| for vad_turn in vad_output.get_timeline().support(): | |
| vad_start = _start + int(sampling_rate * vad_turn.start) | |
| vad_end = _start + int(sampling_rate * vad_turn.end) | |
| prediction = pipe(waveform[vad_start: vad_end])['text'] | |
| history += f"{counter}\n" + \ | |
| f"{second_to_timecode(speech_turn.start + vad_turn.start)} --> {second_to_timecode(speech_turn.start + vad_turn.end)}\n" + \ | |
| f"{prediction}\n\n" | |
| # f">> {speaker}: {prediction}\n\n" | |
| yield history, history, None | |
| counter += 1 | |
| else: | |
| prediction = pipe(data)['text'] | |
| history += f"{counter}\n" + \ | |
| f"{second_to_timecode(speech_turn.start)} --> {second_to_timecode(speech_turn.end)}\n" + \ | |
| f"{prediction}\n\n" | |
| # f">> {speaker}: {prediction}\n\n" | |
| counter += 1 | |
| yield history, history, None | |
| # https://support.google.com/youtube/answer/2734698?hl=en#zippy=%2Cbasic-file-formats%2Csubrip-srt-example%2Csubviewer-sbv-example | |
| file_name = 'transcript.srt' | |
| with open(file_name, 'w') as fp: | |
| fp.write(history) | |
| yield history, history, file_name | |
| demo = gr.Interface( | |
| generator, | |
| inputs=[ | |
| gr.inputs.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL", optional=True), | |
| gr.inputs.Audio(source="microphone", type="filepath", optional=True), | |
| gr.inputs.Audio(source="upload", type="filepath", optional=True), | |
| gr.Number(value=1, label="Number of Speakers"), | |
| gr.Number(value=120, label="Maximum Duration (Seconds)"), | |
| 'state', | |
| ], | |
| outputs=['text', 'state', 'file'], | |
| layout="horizontal", | |
| theme="huggingface", | |
| title="Transcribe Mongolian Whisper 🇲🇳", | |
| description=( | |
| "Transcribe Youtube Video / Microphone / Uploaded File in Mongolian Whisper Model." + \ | |
| " | You can upload SubRip file (`.srt`) [to your youtube video](https://support.google.com/youtube/answer/2734698?hl=en#zippy=%2Cbasic-file-formats)." + \ | |
| " | Please REFRESH 🔄 the page after you transcribed!" + \ | |
| " | 🐦 [@_tsogoo_](https://twitter.com/_tsogoo_)" + \ | |
| " | 🤗 [@bayartsogt](https://huggingface.co/bayartsogt)" + \ | |
| "" | |
| ), | |
| allow_flagging="never", | |
| ) | |
| # define queue - required for generators | |
| demo.queue() | |
| demo.launch() |