Spaces:
Runtime error
Runtime error
import os | |
import torch | |
import gradio as gr | |
import time | |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline | |
from flores200_codes import flores_codes | |
import vinai_translator | |
def load_models(): | |
# build model and tokenizer | |
model_name_dict = {#'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M', | |
#'nllb-1.3B': 'facebook/nllb-200-1.3B', | |
'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B', | |
#'nllb-3.3B': 'facebook/nllb-200-3.3B', | |
} | |
model_dict = {} | |
for call_name, real_name in model_name_dict.items(): | |
print('\tLoading model: %s' % call_name) | |
model = AutoModelForSeq2SeqLM.from_pretrained(real_name) | |
tokenizer = AutoTokenizer.from_pretrained(real_name) | |
model_dict[call_name+'_model'] = model | |
model_dict[call_name+'_tokenizer'] = tokenizer | |
return model_dict | |
def translation(source, target, text): | |
source = flores_codes[source] | |
target = flores_codes[target] | |
output = "" | |
start_time = time.time() | |
chunks = text.splitlines(True) | |
if source == "vie_Latn" or target == "eng_Latn": | |
for chunk in chunks: | |
stchunk = vinai_translator.translate_vi2en(chunk) | |
output += stchunk+"\n" | |
elif source == "eng_Latn" or target == "vie_Latn": | |
for chunk in chunks: | |
stchunk = vinai_translator.translate_en2vi(chunk) | |
output += stchunk+"\n" | |
else: | |
if len(model_dict) == 2: | |
#model_name = 'nllb-distilled-600M' | |
model_name = 'nllb-distilled-1.3B' | |
model = model_dict[model_name + '_model'] | |
tokenizer = model_dict[model_name + '_tokenizer'] | |
translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source, tgt_lang=target) | |
for chunk in chunks: | |
stchunk = translator(chunk+"<mask>", max_length=400, num_beams=5) | |
output += stchunk[0]['translation_text']+"\n" | |
end_time = time.time() | |
result = {'inference_time': end_time - start_time, | |
'source': source, | |
'target': target, | |
'result': output} | |
return result | |
if __name__ == '__main__': | |
print('\tinit models') | |
global model_dict | |
model_dict = load_models() | |
# define gradio demo | |
lang_codes = list(flores_codes.keys()) | |
#inputs = [gr.inputs.Radio(['nllb-distilled-600M', 'nllb-1.3B', 'nllb-distilled-1.3B'], label='NLLB Model'), | |
inputs = [gr.inputs.Dropdown(lang_codes, default='English', label='Source'), | |
gr.inputs.Dropdown(lang_codes, default='Korean', label='Target'), | |
gr.inputs.Textbox(lines=5, label="Input text"), | |
] | |
outputs = gr.outputs.JSON() | |
title = "NLLB distilled 600M demo" | |
demo_status = "Demo is running on CPU" | |
description = f"Details: https://github.com/facebookresearch/fairseq/tree/nllb. {demo_status}" | |
examples = [ | |
['English', 'Korean', 'Hi. nice to meet you'] | |
] | |
gr.Interface(translation, | |
inputs, | |
outputs, | |
title=title, | |
description=description, | |
).launch() |