| import gradio as gr |
| import yaml |
| import torch |
| from mmtafrica import load_params, translate |
| from huggingface_hub import hf_hub_download |
|
|
| language_map = {'English':'en','Swahili':'sw','Fon':'fon','Igbo':'ig', |
| 'Kinyarwanda':'rw','Xhosa':'xh','Yoruba':'yo','French':'fr'} |
| |
| available_languages = list(language_map.keys()) |
|
|
| |
| checkpoint = hf_hub_download(repo_id="chrisjay/mmtafrica", filename="mmt_translation.pt") |
| device = 'gpu' if torch.cuda.is_available() else 'cpu' |
| params = load_params({'checkpoint':checkpoint,'device':device}) |
|
|
|
|
| def get_translation(source_language,target_language,source_sentence=None): |
| ''' |
| This takes a sentence and gets the translation. |
| ''' |
| |
| source_language_ = language_map[source_language] |
| target_language_ = language_map[target_language] |
| |
| |
| try: |
| pred = translate(params,source_sentence,source_lang=source_language_,target_lang=target_language_) |
| if pred=='': |
| return f"Could not find translation" |
| else: |
| return pred |
| except Exception as error: |
| return f"Issue with translation: \n {error}" |
| |
| title = "MMTAfrica: Multilingual Machine Translation" |
| description = "Enjoy our MMT model that allows you to translate among 6 African languages, English and French!\n\nProfitez de notre modèle MMT qui vous permet de traduire parmi 6 langues africaines, anglais et français!" |
|
|
| iface = gr.Interface(fn=get_translation, |
| inputs=[gr.inputs.Dropdown(choices = available_languages,default='Igbo'), |
| gr.inputs.Dropdown(choices = available_languages,default='Fon'), |
| gr.inputs.Textbox(label="Input")], |
| outputs=gr.outputs.Textbox(type="auto", label='Translation'), |
| title=title, |
| description=description, |
| enable_queue=True, |
| theme='huggingface') |
| iface.launch() |
