Merge pull request #39 from 01-ai/reed/add_eos_token

feat: use '\n' as default eos_token in demo

demo/README.md

@@ -7,6 +7,7 @@ python text_generation.py \
     --model 01-ai/Yi-6B \
     --tokenizer 01-ai/Yi-6B \
     --max-tokens 512 \
+    --eos-token $'\n' \
     --streaming
 ```
 
@@ -21,5 +22,7 @@ torchrun --nproc_per_node 2 \
     text_generation_tp.py \
     --model 01-ai/Yi-6B \
     --max-tokens 512 \
+    --eos-token $'\n' \
     --streaming
+
 ```

demo/text_generation.py

@@ -34,6 +34,12 @@ def parse_inputs():
         default="Let me tell you an interesting story about cat Tom and mouse Jerry,",
         help="The prompt to start with",
     )
+    parser.add_argument(
+        "--eos-token",
+        type=str,
+        default="<|endoftext|>",
+        help="End of sentence token",
+    )
     args = parser.parse_args()
     return args
 
@@ -55,6 +61,7 @@ def main(args):
         inputs.input_ids.cuda(),
         max_new_tokens=args.max_tokens,
         streamer=streamer,
+        eos_token_id=tokenizer.convert_tokens_to_ids(args.eos_token),
         do_sample=True,
     )
     if streamer is None:

demo/text_generation_tp.py

@@ -41,6 +41,12 @@ def parse_inputs():
         default="Let me tell you an interesting story about cat Tom and mouse Jerry,",
         help="The prompt to start with",
     )
+    parser.add_argument(
+        "--eos-token",
+        type=str,
+        default="<|endoftext|>",
+        help="End of sentence token",
+    )
     args = parser.parse_args()
     return args
 
@@ -93,6 +99,7 @@ def on_finalized_text(self, text: str, stream_end: bool = False):
         inputs.input_ids.cuda(),
         max_new_tokens=args.max_tokens,
         streamer=streamer,
+        eos_token_id=tokenizer.convert_tokens_to_ids(args.eos_token),
         do_sample=True,
     )
     if distributed.get_rank() == 0 and streamer is None: