Update README.md

README.md

@@ -86,6 +86,9 @@ if input_ids != None:
 
 print("Updated max_len = " + str(max_len))
 
+stop_token = "<|endoftext|>"
+new_lines = "\n\n\n"
+
 sample_outputs = model.generate(
     input_ids,
     do_sample=True, 
@@ -95,9 +98,18 @@ sample_outputs = model.generate(
     num_return_sequences=sample_output_num
 )
 
-print(100 * '-' + "\nOutput:\n" + 100 * '-')
+print(100 * '-' + "\n\t\tOutput\n" + 100 * '-')
 for i, sample_output in enumerate(sample_outputs):
-  print("\n{}: {}".format(i, tokenizer.decode(sample_output, skip_special_tokens=True)))
+
+  text = tokenizer.decode(sample_output, skip_special_tokens=True)
+  
+  # Remove all text after the stop token
+  text = text[: text.find(stop_token) if stop_token else None]
+
+  # Remove all text after 3 newlines
+  text = text[: text.find(new_lines) if new_lines else None]
+
+  print("\n{}: {}".format(i, text))
   print("\n" + 100 * '-')
 
 ```