Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| import torch | |
| 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): | |
| 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, | |
| "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 "\n".join(output) | |
| with gr.Blocks(theme="gstaff/xkcd") as app: | |
| markdown = """ | |
| # ante toki | |
| ## toki a! 💬 | |
| This is an English to Toki Pona / Toki Pona to English translator. | |
| Input your text to translate and desired number of generated sentences ! | |
| """ | |
| with gr.Row(): | |
| gr.Markdown(markdown) | |
| with gr.Column(): | |
| input_text = gr.components.Textbox( | |
| label="Input Text", | |
| value="Dogs are cute, but I prefer cats.", | |
| ) | |
| source_lang = gr.components.Dropdown( | |
| label="Source Language", | |
| value="English", | |
| choices=list(LANG_CODES.keys()), | |
| ) | |
| target_lang = gr.components.Dropdown( | |
| label="Target Language", | |
| value="Toki Pona", | |
| choices=list(LANG_CODES.keys()), | |
| ) | |
| return_seqs = gr.Slider( | |
| label="Number of generated sentences", | |
| value=3, | |
| minimum=1, | |
| maximum=12, | |
| step=1, | |
| ) | |
| inputs = [input_text, source_lang, target_lang, return_seqs] | |
| outputs = gr.Textbox() | |
| translate_btn = gr.Button("o ante toki!") | |
| translate_btn.click(translate, inputs=inputs, outputs=outputs) | |
| app.launch() | |