Spaces:
Running
Running
File size: 2,186 Bytes
d28a390 1e0bf7f d28a390 1e0bf7f d28a390 1e0bf7f 4aa6835 1e0bf7f 513bb2f 1e0bf7f d28a390 1e0bf7f d28a390 1e0bf7f d28a390 65f9b13 fd42785 |
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 |
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from flores200_codes import flores_codes
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, model_name="nllb-distilled-600M"):
if len(model_dict) == 2:
model_name = 'nllb-distilled-600M'
start_time = time.time()
source = flores_codes[source]
target = flores_codes[target]
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)
output = translator(text, max_length=400)
end_time = time.time()
output = output[0]['translation_text']
result = {'inference_time': end_time - start_time,
'source': source,
'target': target,
'result': output}
return result
if __name__ == '__main__':
model_dict = load_models()
lang_codes = list(flores_codes.keys())
inputs = [
gr.components.Dropdown(lang_codes, label='Source'),
gr.components.Dropdown(lang_codes, label='Target'),
gr.components.Textbox(lines=5, label="Input text"),
gr.components.Dropdown(["nllb-distilled-600M"], label="Model"),
]
outputs = gr.components.JSON()
title = "NLLB distilled 600M demo"
demo_status = "Demo is running on CPU"
gr.Interface(translation, inputs, outputs, title=title).launch(server_port=450)
|