|
import gradio as gr |
|
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline |
|
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): |
|
""" |
|
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, |
|
'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() as app: |
|
markdown=""" |
|
# An English / toki pona Neural Machine Translation App! |
|
|
|
### toki a! 💬 |
|
This is an english to toki pona / toki pona to english neural machine translation app. |
|
Input your text to translate, a source language and target language, and desired number of return sequences! |
|
### Grammar Regularization |
|
An interesting quirk of training a many-to-many translation model is that pseudo-grammar correction |
|
can be achieved by translating *from* **language A** *to* **language A** |
|
|
|
Remember, this can ***approximate*** grammaticality, but it isn't always the best. |
|
|
|
For example, "mi li toki e toki pona" (Source Language: toki pona & Target Language: toki pona) will result in: |
|
- ['mi toki e toki pona.', 'mi toki pona.', 'mi toki e toki pona'] |
|
- (Thus, the ungrammatical "li" is dropped) |
|
### Model and Data |
|
This app utilizes a fine-tuned version of Facebook/Meta AI's M2M100 418M param model. |
|
|
|
By leveraging the pretrained weights of the massively multilingual M2M100 model, |
|
we can jumpstart our transfer learning to accomplish machine translation for toki pona! |
|
|
|
The model was fine-tuned on the English/toki pona bitexts found at [https://tatoeba.org/](https://tatoeba.org/) |
|
|
|
### This app is a work in progress and obviously not all translations will be perfect. |
|
In addition to parameter quantity and the hyper-parameters used while training, |
|
the *quality of data* found on Tatoeba directly influences the perfomance of projects like this! |
|
If you wish to contribute, please add high quality and diverse translations to Tatoeba! |
|
""" |
|
|
|
with gr.Row(): |
|
gr.Markdown(markdown) |
|
with gr.Column(): |
|
input_text = gr.components.Textbox(label="Input Text", value="Raccoons are fascinating creatures, but I prefer opossums.") |
|
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 return sequences", value=3, minimum=1, maximum=12, step=1) |
|
|
|
inputs=[input_text, source_lang, target_lang, return_seqs] |
|
outputs = gr.Textbox() |
|
|
|
translate_btn = gr.Button("Translate! | o ante toki!") |
|
translate_btn.click(translate, inputs=inputs, outputs=outputs) |
|
|
|
gr.Examples( |
|
[ |
|
["Hello! How are you?", "English", "toki pona", 3], |
|
["toki a! ilo pi ante toki ni li pona!", "toki pona", "English", 3], |
|
["mi li toki e toki pona", "toki pona", "toki pona", 3], |
|
], |
|
inputs=inputs |
|
) |
|
|
|
app.launch() |
