Spaces:
Sleeping
Sleeping
File size: 2,138 Bytes
0ad5b35 d299dc7 0ad5b35 f3b9070 2bb196b 0ad5b35 9c1e6c1 0ad5b35 4d14df0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import gradio as gr
from transformers import pipeline
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
from utils import lang_ids
import nltk
nltk.download('punkt')
MODEL_NAME = "Pranjal12345/pranjal_whisper_medium"
BATCH_SIZE = 10
FILE_LIMIT_MB = 1000
pipe = pipeline(
task="automatic-speech-recognition",
model=MODEL_NAME,
chunk_length_s=30,
device='cpu',
)
## Download the mbart model from hugging face
model = MBartForConditionalGeneration.from_pretrained("sanjitaa/mbart-many-to-many")
tokenizer = MBart50TokenizerFast.from_pretrained("sanjitaa/mbart-many-to-many")
lang_list = list(lang_ids.keys())
def translate_audio(inputs,target_language):
if inputs is None:
raise gr.Error("No audio file submitted! Please upload an audio file before submitting your request.")
text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": "translate"}, return_timestamps=True)["text"]
target_lang = lang_ids[target_language]
if target_language == 'English':
return text
else:
tokenizer.src_lang = "en_XX"
chunks = nltk.tokenize.sent_tokenize(text)
translated_text = ''
for segment in chunks:
encoded_chunk = tokenizer(segment, return_tensors="pt")
generated_tokens = model.generate(
**encoded_chunk,
forced_bos_token_id=tokenizer.lang_code_to_id[target_lang]
)
translated_chunk = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
translated_text = translated_text + translated_chunk[0]
return translated_text
inputs=[
gr.inputs.Audio(source = "upload", type="filepath", label="Audio file"),
gr.Dropdown(lang_list, value="English", label="Target Language"),
]
description = "Audio translation"
translation_interface = gr.Interface(
fn=translate_audio,
inputs= inputs,
outputs="text",
title="Speech Translation",
description= description
)
if __name__ == "__main__":
translation_interface.launch()
|