Updated app.py

app.py

@@ -1,13 +1,34 @@
 import gradio as gr
-from simpletransformers.seq2seq import Seq2SeqModel
+from simpletransformers.seq2seq import Seq2SeqModel, Seq2SeqArgs
 
 # Define the models' paths
 BM_MODEL_PATH = "Enutrof/marian-mt-en-pcm"
 BBGM_MODEL_PATH = "NITHUB-AI/marian-mt-bbc-en-pcm"
 
+def load_translator(model_name='Enutrof/marian-mt-en-pcm'):
+    '''
+    This method loads the sequence to sequence model for translation.
+    :return: model
+    '''
+    pmodel_args = Seq2SeqArgs()
+    pmodel_args.max_length = 1024
+    pmodel_args.length_penalty = 1
+    pmodel_args.num_beams = 20
+    pmodel_args.num_return_sequences = 3
+
+    pmodel = Seq2SeqModel(
+        encoder_decoder_type="marian",
+        encoder_decoder_name=model_name,
+        args=pmodel_args,
+        use_cuda=False
+    )
+    return pmodel
+
 #Load models
-bm_model = Seq2SeqModel(encoder_decoder_type="marian", encoder_decoder_name=BM_MODEL_PATH, use_cuda=False)
-bbgm_model = Seq2SeqModel(encoder_decoder_type="marian", encoder_decoder_name=BBGM_MODEL_PATH, use_cuda=False)
+
+bm_model = load_translator(BM_MODEL_PATH)
+bbgm_model = load_translator(BBGM_MODEL_PATH)
+
 
 # Dictionary to easily select model
 models = {
@@ -16,9 +37,11 @@ models = {
 }
 
 def translate(model_name, source_sentence, num_beams):
-    selected_model = models[model_name]
-    predictions = selected_model.predict([source_sentence] * 3, num_beams=int(num_beams), num_return_sequences=3)
-    return tuple(predictions)
+    if isinstance(source_sentence, str):
+        source_sentence = [source_sentence]
+    model = models[model_name]
+    predictions = model.predict(input, num_beams=int(num_beams))
+    return [i.replace('▁', ' ') for i in predictions[0]]
 
 # Gradio interface
 interface = gr.Interface(
@@ -33,7 +56,17 @@ interface = gr.Interface(
         gr.Textbox(label="Prediction 2"),
         gr.Textbox(label="Prediction 3"),
     ],
-    live=True
+    title='English to 🇳🇬 Pidgin Automatic Translation',
+    description='Type your English text in the left text box to get 🇳🇬 Pidgin translations on the right. '
+                'Tell us the best translation by clicking one of the buttons below.',
+    examples=[
+        'Who are you?',
+        'Is a personal philosophy of moral relativism, the only way to survive in this ethically complex world, or is it just an excuse to justify doing bad things?',
+        'I know every song by that artiste.',
+        'They should not be permitted here.',
+        'What are you looking for?',
+        'I am lost, please help me find my way to the market.',
+    ]
 )
 
-interface.launch()
+interface.launch(enable_queue=True)