Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline | |
| import torch | |
| LANGS = ["kin_Latn","eng_Latn"] | |
| TASK = "translation" | |
| # CKPT = "DigitalUmuganda/Finetuned-NLLB" | |
| MODELS = ["facebook/nllb-200-distilled-600M","DigitalUmuganda/Finetuned-NLLB"] | |
| # model = AutoModelForSeq2SeqLM.from_pretrained(CKPT) | |
| # tokenizer = AutoTokenizer.from_pretrained(CKPT) | |
| device = 0 if torch.cuda.is_available() else -1 | |
| fb_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M") | |
| du_model = AutoModelForSeq2SeqLM.from_pretrained("DigitalUmuganda/Finetuned-NLLB") | |
| tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M") | |
| models = {"facebook/nllb-200-distilled-600M":fb_model,"DigitalUmuganda/Finetuned-NLLB":du_model} | |
| # def translate(text, src_lang, tgt_lang, max_length=400): | |
| def translate_fb(text, src_lang, tgt_lang, max_length=400): | |
| """ | |
| Translate the text from source lang to target lang | |
| """ | |
| print("fb src_lang: ",src_lang) | |
| print("fb dest_lang: ",tgt_lang) | |
| translation_pipeline = pipeline(TASK, | |
| tokenizer=tokenizer, | |
| src_lang=src_lang, | |
| tgt_lang=tgt_lang, | |
| model = fb_model, | |
| max_length=max_length, | |
| device=device) | |
| result = translation_pipeline(text) | |
| return result[0]['translation_text'] | |
| def translate_du(text, src_lang, tgt_lang, CKPT, max_length=400): | |
| """ | |
| Translate the text from source lang to target lang | |
| """ | |
| print("du src_lang: ",src_lang) | |
| print("du tgt_lang: ",tgt_lang) | |
| translation_pipeline = pipeline(TASK, | |
| tokenizer=tokenizer, | |
| src_lang=src_lang, | |
| tgt_lang=tgt_lang, | |
| model = du_model, | |
| max_length=max_length, | |
| device=device) | |
| result = translation_pipeline(text) | |
| return result[0]['translation_text'] | |
| gr_fb = gr.Interface( | |
| translate_fb, | |
| [ | |
| 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, 512, value=400, step=8, label="Max Length") | |
| ], | |
| ['text'], | |
| #examples=examples, | |
| # article=article, | |
| cache_examples=False, | |
| title="nllb-200-distilled-600M", | |
| #description=description | |
| ) | |
| gr_du = gr.Interface( | |
| translate_du, | |
| [ | |
| 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, 512, value=400, step=8, label="Max Length") | |
| ], | |
| ['text'], | |
| #examples=examples, | |
| # article=article, | |
| cache_examples=False, | |
| title="nllb-200-distilled-600M-Finetuned", | |
| # description=description | |
| ) | |
| gr.Parallel( | |
| gr_fb, | |
| gr_du, | |
| # [ | |
| # 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, 512, value=400, step=8, label="Max Length") | |
| ).launch() | |