Update app.py

app.py

@@ -60,11 +60,15 @@ def translate_vi2en(vi_text: str) -> str:
     input_ids = tokenizer_vi2en(vi_text, return_tensors="pt").input_ids
     output_ids = model_vi2en.generate(
         input_ids,
-        do_sample=True,
-        top_k=100,
-        top_p=0.8,
         decoder_start_token_id=tokenizer_vi2en.lang_code_to_id["en_XX"],
         num_return_sequences=1,
+        # # With sampling
+        # do_sample=True,
+        # top_k=100,
+        # top_p=0.8,
+        # With beam search
+        num_beams=5,
+        early_stopping=True
     )
     en_text = tokenizer_vi2en.batch_decode(output_ids, skip_special_tokens=True)
     en_text = " ".join(en_text)
@@ -77,11 +81,15 @@ def translate_en2vi(en_text: str) -> str:
     input_ids = tokenizer_en2vi(en_text, return_tensors="pt").input_ids
     output_ids = model_en2vi.generate(
         input_ids,
-        do_sample=True,
-        top_k=100,
-        top_p=0.8,
         decoder_start_token_id=tokenizer_en2vi.lang_code_to_id["vi_VN"],
         num_return_sequences=1,
+        # # With sampling
+        # do_sample=True,
+        # top_k=100,
+        # top_p=0.8,
+        # With beam search
+        num_beams=5,
+        early_stopping=True
     )
     vi_text = tokenizer_en2vi.batch_decode(output_ids, skip_special_tokens=True)
     vi_text = " ".join(vi_text)