| """Pronkin_hw_task3.ipynb | |
| https://colab.research.google.com/drive/149j9u-wsD3GiEwRA8clBrXQ8bh5DRk7I?usp=sharing | |
| """ | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from datasets import load_dataset | |
| from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor, pipeline, WhisperProcessor, VitsModel, VitsTokenizer | |
| device = "cuda:0" if torch.cuda.is_available() else "cpu" | |
| asr_pipe = pipeline("automatic-speech-recognition", model="voidful/wav2vec2-xlsr-multilingual-56", device=device) | |
| processor = WhisperProcessor.from_pretrained("openai/whisper-small") | |
| translator_1 = pipeline("translation", model="Helsinki-NLP/opus-mt-mul-en") | |
| translator_2 = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ru") | |
| model = VitsModel.from_pretrained("facebook/mms-tts-rus") | |
| tokenizer = VitsTokenizer.from_pretrained("facebook/mms-tts-rus") | |
| def translator_mul_ru(text): | |
| translation = translator_2(translator_1(text)[0]['translation_text']) | |
| return translation[0]['translation_text'] | |
| def translate(audio): | |
| outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate"}) | |
| return outputs["text"] | |
| def synthesise(text): | |
| translated_text = translator_mul_ru(text) | |
| inputs = tokenizer(translated_text, return_tensors="pt") | |
| input_ids = inputs["input_ids"] | |
| with torch.no_grad(): | |
| outputs = model(input_ids) | |
| speech = outputs["waveform"] | |
| return speech.cpu() | |
| def speech_to_speech_translation(audio): | |
| translated_text = translate(audio) | |
| print(translated_text) | |
| synthesised_speech = synthesise(translated_text) | |
| synthesised_speech = (synthesised_speech.numpy() * 32767).astype(np.int16) | |
| return 16000, synthesised_speech[0] | |
| title = "Pronkin custom STST" | |
| description = """ | |
| * ASR-модель распознает речь с помощью voidful/wav2vec2-xlsr-multilingual-56 и возвращает текст на любом из 56 языков. | |
| * Перевод текста с любого на английский с помощью модели Helsinki-NLP/opus-mt-mul-en, с английского на русский - Helsinki-NLP/opus-mt-en-ru | |
| * Синтез речи на русском языке с помощью модели facebook/mms-tts-rus | |
| """ | |
| demo = gr.Blocks() | |
| mic_translate = gr.Interface( | |
| fn=speech_to_speech_translation, | |
| inputs=gr.Audio(source="microphone", type="filepath"), | |
| outputs=gr.Audio(label="Generated Speech", type="numpy"), | |
| title=title, | |
| description=description, | |
| ) | |
| file_translate = gr.Interface( | |
| fn=speech_to_speech_translation, | |
| inputs=gr.Audio(source="upload", type="filepath"), | |
| outputs=gr.Audio(label="Generated Speech", type="numpy"), | |
| title=title, | |
| description=description, | |
| ) | |
| with demo: | |
| gr.TabbedInterface([mic_translate, file_translate], ["Microphone", "File"]) | |
| demo.launch() | |