Update app.py

app.py

@@ -5,10 +5,9 @@ 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"
-}
+model.to(device)
+LANG_CODES = {"English": "en", "toki pona": "tl"}
+
 
 def translate(text):
     """
@@ -20,51 +19,47 @@ def translate(text):
 
     tokenizer.src_lang = "en"
     tokenizer.tgt_lang = "tl"
-
-    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': 1,
-            'num_beams':1,
-            'forced_bos_token_id': tokenizer.lang_code_to_id["tl"]
-        }
-    
+        "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": 1,
+        "num_beams": 1,
+        "forced_bos_token_id": tokenizer.lang_code_to_id["tl"],
+    }
 
     outs = model.generate(**{**ins, **gen_args})
     output = tokenizer.batch_decode(outs.sequences, skip_special_tokens=True)
-    text2 = '\n'.join(output)
-    
+    text2 = "\n".join(output)
+
     ##################
 
     tokenizer.src_lang = "tl"
     tokenizer.tgt_lang = "en"
 
-    ins = tokenizer(text2, return_tensors='pt').to(device)
+    ins = tokenizer(text2, 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': 1,
-            'num_beams':1,
-            'forced_bos_token_id': tokenizer.lang_code_to_id["en"]
-        }
-    
+        "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": 1,
+        "num_beams": 1,
+        "forced_bos_token_id": tokenizer.lang_code_to_id["en"],
+    }
 
     outs2 = model.generate(**{**ins, **gen_args})
     output2 = tokenizer.batch_decode(outs2.sequences, skip_special_tokens=True)
 
+    return "\n".join(output2)
 
 
-    return '\n'.join(output2)
-
 with gr.Blocks() as app:
-    markdown="""
+    markdown = """
     # An English / toki pona Neural Machine Translation App!
     
     ### toki a! 💬
@@ -101,12 +96,15 @@ with gr.Blocks() as app:
     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.")
+            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]
+
+            inputs = [input_text]
             outputs = gr.Textbox()
 
             translate_btn = gr.Button("Translate! | o ante toki!")
@@ -115,10 +113,10 @@ with gr.Blocks() as app:
             gr.Examples(
                 [
                     ["Hello! How are you?", "English", "toki pona", 3],
-                    ["toki a! ilo pi ante toki ni li pona!", "toki pona", "English",  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
+                inputs=inputs,
             )
 
-app.launch()
+app.launch()