import gradio as gr from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline import torch # this model was loaded from https://hf.co/models model = AutoModelForSeq2SeqLM.from_pretrained("Jayyydyyy/m2m100_418m_tokipona") tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M") device = "cuda:0" if torch.cuda.is_available() else "cpu" LANG_CODES = { "English":"en", "toki pona":"tl" } def translate(text, src_lang, tgt_lang, candidates:int): """ Translate the text from source lang to target lang """ src = LANG_CODES.get(src_lang) tgt = LANG_CODES.get(tgt_lang) tokenizer.src_lang = src tokenizer.tgt_lang = tgt ins = tokenizer(text, return_tensors='pt').to(device) gen_args = { 'return_dict_in_generate': True, 'output_scores': True, 'output_hidden_states': True, 'length_penalty': 0.0, # don't encourage longer or shorter output, 'num_return_sequences': candidates, 'num_beams':candidates, 'forced_bos_token_id': tokenizer.lang_code_to_id[tgt] } outs = model.generate(**{**ins, **gen_args}) output = tokenizer.batch_decode(outs.sequences, skip_special_tokens=True) return output app = gr.Interface( fn=translate, inputs=[ gr.components.Textbox(label="Text"), gr.components.Dropdown(label="Source Language", choices=list(LANG_CODES.keys())), gr.components.Dropdown(label="Target Language", choices=list(LANG_CODES.keys())), gr.Slider(label="Number of return sequences", value=3, minimum=1, maximum=12, step=1) ], outputs=["text"], examples=[ ["Welcome to my translation app.", "English", "toki pona", 3], ["Its not always perfect, but its pretty okay!", "English", "toki pona", 3], ["ilo pi ante toki ni li pona a!", "toki pona", "English", 3], ["kijetesantakalu li pona", "toki pona", "English", 3], ], cache_examples=False, title="A simple English / toki pona Neural Translation App", description="A simple English / toki pona Neural Translation App" ) app.launch()