Norod78 commited on
Commit
8f131b4
1 Parent(s): fca5f65

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +15 -7
README.md CHANGED
@@ -85,6 +85,9 @@ if input_ids != None:
85
 
86
  print("Updated max_len = " + str(max_len))
87
 
 
 
 
88
  sample_outputs = model.generate(
89
  input_ids,
90
  do_sample=True,
@@ -94,13 +97,18 @@ sample_outputs = model.generate(
94
  num_return_sequences=sample_output_num
95
  )
96
 
97
- print(100 * '-' + "\
98
- Output:\
99
- " + 100 * '-')
100
  for i, sample_output in enumerate(sample_outputs):
101
- print("\
102
- {}: {}".format(i, tokenizer.decode(sample_output, skip_special_tokens=True)))
103
- print("\
104
- " + 100 * '-')
 
 
 
 
 
 
 
105
 
106
  ```
 
85
 
86
  print("Updated max_len = " + str(max_len))
87
 
88
+ stop_token = "<|endoftext|>"
89
+ new_lines = "\n\n\n"
90
+
91
  sample_outputs = model.generate(
92
  input_ids,
93
  do_sample=True,
 
97
  num_return_sequences=sample_output_num
98
  )
99
 
100
+ print(100 * '-' + "\n\t\tOutput\n" + 100 * '-')
 
 
101
  for i, sample_output in enumerate(sample_outputs):
102
+
103
+ text = tokenizer.decode(sample_output, skip_special_tokens=True)
104
+
105
+ # Remove all text after the stop token
106
+ text = text[: text.find(stop_token) if stop_token else None]
107
+
108
+ # Remove all text after 3 newlines
109
+ text = text[: text.find(new_lines) if new_lines else None]
110
+
111
+ print("\n{}: {}".format(i, text))
112
+ print("\n" + 100 * '-')
113
 
114
  ```