Spaces:
Runtime error
Runtime error
File size: 1,991 Bytes
666c637 d446722 666c637 b7fade1 666c637 a7ce701 666c637 6302ad4 666c637 |
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 |
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() |