Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline | |
| import torch | |
| TASK = "translation" | |
| MODEL = "facebook/nllb-200-distilled-600M" | |
| LANGS = ["Kinyarwanda","English","French","Swahili","Luganda","Lingala"] | |
| LANGUAGES={ | |
| "english": "eng_Latn", | |
| "kinyarwanda": "kin_Latn", | |
| "swahili": "swh_Latn", | |
| "luganda": "lug_Latn", | |
| "lingala": "lin_Latn", | |
| "french": "fra_Latn" | |
| } | |
| device = 0 if torch.cuda.is_available() else -1 | |
| model = AutoModelForSeq2SeqLM.from_pretrained(MODEL) | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL) | |
| def choose_language(language): | |
| result = LANGUAGES.get(language.lower()) | |
| if result: | |
| return result | |
| return | |
| def translate(text, source_lang, target_lang, max_length=400): | |
| """ | |
| Translate text from source language to target language | |
| """ | |
| src_lang = choose_language(source_lang) | |
| tgt_lang= choose_language(target_lang) | |
| if src_lang==None: | |
| return "Error: the source langage is incorrect" | |
| elif tgt_lang==None: | |
| return "Error: the target language is incorrect" | |
| translation_pipeline = pipeline(TASK, | |
| model=model, | |
| tokenizer=tokenizer, | |
| src_lang=src_lang, | |
| tgt_lang=tgt_lang, | |
| max_length=max_length, | |
| device=device) | |
| result = translation_pipeline(text) | |
| return result[0]['translation_text'] | |
| gradio_ui= gr.Interface( | |
| fn=translate, | |
| title="NLLB 600M Translation Demo", | |
| inputs= [ | |
| gr.components.Textbox(label="Text"), | |
| gr.components.Dropdown(label="Source Language", choices=LANGS), | |
| gr.components.Dropdown(label="Target Language", choices=LANGS), | |
| gr.components.Slider(8, 400, value=400, step=8, label="Max Length") | |
| ], | |
| outputs=gr.outputs.Textbox(label="Translated text") | |
| ) | |
| gradio_ui.launch() |