audio-translate / app.py
vumichien's picture
Update app.py
2347c54
import gradio as gr
import librosa
from optimum.onnxruntime import ORTModelForSeq2SeqLM
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
# load model and processor
processor = Wav2Vec2Processor.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-english")
model = Wav2Vec2ForCTC.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-english")
tokenizer = AutoTokenizer.from_pretrained("icon-it-tdtu/mt-en-vi-optimum")
model_lm = ORTModelForSeq2SeqLM.from_pretrained("icon-it-tdtu/mt-en-vi-optimum")
def process_audio_file(file):
data, sr = librosa.load(file)
if sr != 16000:
data = librosa.resample(data, sr, 16000)
print(data.shape)
inputs = processor(data, sampling_rate=16000, return_tensors="pt", padding=True)
return inputs
def interpret(micro, file):
if file is not None and micro is None:
input_audio = file
elif file is None and micro is not None:
input_audio = micro
else:
input_audio = file
inputs = process_audio_file(input_audio)
with torch.no_grad():
output_logit = model(inputs.input_values).logits
pred_ids = torch.argmax(output_logit, dim=-1)
text = processor.batch_decode(pred_ids)[0].lower()
print(text)
translate_text = translate(text)
return translate_text
def translate(text):
batch = tokenizer([text], return_tensors="pt")
generated_ids = model_lm.generate(**batch)
translated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
return translated_text
# Set the starting state to an empty string
iface = gr.Interface(
fn=interpret,
title="Interpret English to Vietnamese",
description="A simple interface to interpret from spoken English to Vietnamese.",
article="Author: <a href=\"https://huggingface.co/vumichien\">Vu Minh Chien</a>.",
inputs=[
gr.Audio(source="microphone", type="filepath", streaming=False, label="Microphone"),
gr.Audio(source="upload", type="filepath", label="File"),
],
outputs=gr.Textbox(label="Text output")
,
examples=[[None, "samples/sample-000000.mp3"], [None, "samples/sample-000001.mp3"]],
)
iface.launch(debug=True)