Spaces:
Sleeping
Sleeping
| import os | |
| import torch | |
| import gradio as gr | |
| import time | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline | |
| codes_as_string = '''Assamese asm_Beng | |
| Awadhi awa_Deva | |
| Bengali ben_Beng | |
| Bhojpuri bho_Deva | |
| Standard Tibetan bod_Tibt | |
| Dzongkha dzo_Tibt | |
| English eng_Latn | |
| Gujarati guj_Gujr | |
| Hindi hin_Deva | |
| Chhattisgarhi hne_Deva | |
| Kannada kan_Knda | |
| Kashmiri (Arabic script) kas_Arab | |
| Kashmiri (Devanagari script) kas_Deva | |
| Mizo lus_Latn | |
| Magahi mag_Deva | |
| Maithili mai_Deva | |
| Malayalam mal_Mlym | |
| Marathi mar_Deva | |
| Meitei (Bengali script) mni_Beng | |
| Burmese mya_Mymr | |
| Nepali npi_Deva | |
| Odia ory_Orya | |
| Punjabi pan_Guru | |
| Sanskrit san_Deva | |
| Santali sat_Olck | |
| Sindhi snd_Arab | |
| Tamil tam_Taml | |
| Telugu tel_Telu | |
| Urdu urd_Arab | |
| Vietnamese vie_Latn''' | |
| def load_models(): | |
| # build model and tokenizer | |
| model_name_dict = { | |
| 'nllb-1.3B': "ychenNLP/nllb-200-distilled-1.3B-easyproject", | |
| } | |
| 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("facebook/nllb-200-distilled-600M") | |
| model_dict[call_name+'_model'] = model | |
| model_dict[call_name+'_tokenizer'] = tokenizer | |
| return model_dict | |
| def translation(source, target, text): | |
| if len(model_dict) == 2: | |
| model_name = 'nllb-1.3B' | |
| 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() | |
| full_output = output | |
| output = output[0]['translation_text'] | |
| # result = {'inference_time': end_time - start_time, | |
| # 'source': source, | |
| # 'target': target, | |
| # 'result': output, | |
| # 'full_output': full_output} | |
| return output | |
| if __name__ == '__main__': | |
| print('\tinit models') | |
| codes_as_string = codes_as_string.split('\n') | |
| flores_codes = {} | |
| for code in codes_as_string: | |
| lang, lang_code = code.split('\t') | |
| flores_codes[lang] = lang_code | |
| global model_dict | |
| model_dict = load_models() | |
| # define gradio demo | |
| lang_codes = list(flores_codes.keys()) | |
| inputs = [gr.inputs.Dropdown(lang_codes, default='English', label='Source'), | |
| gr.inputs.Dropdown(lang_codes, default='Hindi', label='Target'), | |
| gr.inputs.Textbox(lines=5, label="Input text"), | |
| ] | |
| outputs = gr.inputs.Textbox(label="Output text") | |
| title = "Machine Translation Demo" | |
| demo_status = "Machine Translation System." | |
| description = f"{demo_status}" | |
| gr.Interface(translation, | |
| inputs, | |
| outputs, | |
| title=title, | |
| description=description, | |
| examples_per_page=50, | |
| theme="JohnSmith9982/small_and_pretty" | |
| ).launch() | |