Update README.md

README.md

@@ -65,7 +65,17 @@ 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
+}
+
 input_text += "<SEP>"
 input_text = preprocess(input_text)
 
@@ -74,9 +84,7 @@ with torch.no_grad():
 
     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)
 output = postprocess(output)