Update app.py

app.py

@@ -1,95 +1,41 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+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"
+	"English":"en",
+	"toki pona":"tl"
 }
 
 def translate(text, src_lang, tgt_lang, candidates:int):
-    """
-    Translate the text from source lang to target lang
-    """
+	"""
+	Translate the text from source lang to target lang
+	"""
 
-    src = LANG_CODES.get(src_lang)
-    tgt = LANG_CODES.get(tgt_lang)
+	src = LANG_CODES.get(src_lang)
+	tgt = LANG_CODES.get(tgt_lang)
 
-    tokenizer.src_lang = src
-    tokenizer.tgt_lang = tgt
+	tokenizer.src_lang = src
+	tokenizer.tgt_lang = tgt
 
-    ins = tokenizer(text, return_tensors='pt').to(device)
+	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]
-        }
-    
+	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 '\n'.join(output)
+	outs = model.generate(**{**ins, **gen_args})
+	output = tokenizer.batch_decode(outs.sequences, skip_special_tokens=True)
 
-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!
-    """
+	return '\n'.join(output)
 
-    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()
+print(translate("Hello!", "English", "toki pona", 1))