Update app.py

app.py

@@ -19,10 +19,23 @@ def generate(input, slang, tlang):
     tlang = CREOLE[tlang]
     inp = tokenizer(input.strip() + " </s> <2" + slang + ">",
                     add_special_tokens=False, return_tensors="pt", padding=True).input_ids
-    model_output = model.generate(inp, use_cache=True, num_beams=1, max_length=int(2*len(inp[0])), min_length=1, early_stopping=True, pad_token_id=pad_id,
+    if (slang != "eng" and tlang == "eng") or (slang == "eng" and tlang != "eng"):
+        model_output = model.generate(inp, use_cache=True, num_beams=1, max_length=int(2*len(inp[0])), min_length=1, early_stopping=True, pad_token_id=pad_id,
                                   bos_token_id=bos_id, eos_token_id=eos_id, decoder_start_token_id=tokenizer._convert_token_to_id_with_added_voc("<2"+tlang+">"))
-    decoded_output = tokenizer.decode(
+        decoded_output = tokenizer.decode(
         model_output[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    elif slang != tlang:
+        model_output = model.generate(inp, use_cache=True, num_beams=1, max_length=int(2*len(inp[0])), min_length=1, early_stopping=True, pad_token_id=pad_id,
+                                  bos_token_id=bos_id, eos_token_id=eos_id, decoder_start_token_id=tokenizer._convert_token_to_id_with_added_voc("<2eng>"))
+        decoded_output = tokenizer.decode(
+        model_output[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        inp = tokenizer(decoded_output + " </s> <2eng>",
+                    add_special_tokens=False, return_tensors="pt", padding=True).input_ids
+        model_output = model.generate(inp, use_cache=True, num_beams=1, max_length=int(2*len(inp[0])), min_length=1, early_stopping=True, pad_token_id=pad_id,
+                                  bos_token_id=bos_id, eos_token_id=eos_id, decoder_start_token_id=tokenizer._convert_token_to_id_with_added_voc("<2"+tlang+">"))
+        decoded_output = tokenizer.decode(
+        model_output[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    
 
     return decoded_output