Update app.py

app.py

@@ -17,6 +17,16 @@ def preprocess(text):
 def postprocess(text):
     return text.replace("<LB>", "\n")
 
+generation_config = {
+    "max_new_tokens": 256,
+    "num_beams": 3,
+    "num_return_sequences": 1,
+    "early_stopping": True,
+    "eos_token_id": tokenizer.eos_token_id,
+    "pad_token_id": tokenizer.pad_token_id,
+    "repetition_penalty": 3.0
+}
+
 def generate(input_text):
     input_text += tokenizer.eos_token
     input_text = preprocess(input_text)
@@ -26,9 +36,7 @@ def generate(input_text):
     
         output_ids = model.generate(
             token_ids.to(model.device),
-            max_new_tokens=256,
-            pad_token_id=tokenizer.pad_token_id,
-            eos_token_id=tokenizer.eos_token_id,
+            **generation_config
         )
     output = tokenizer.decode(output_ids.tolist()[0][token_ids.size(1) :], skip_special_tokens=True)
     return postprocess(output)