update text_generation

Author: Yimi81

demo/text_generation.py

@@ -11,16 +11,10 @@ def parse_inputs():
         default="01-ai/Yi-6B",
         help="pretrained model path locally or name on huggingface",
     )
-    parser.add_argument(
-        "--tokenizer",
-        type=str,
-        default="",
-        help="tokenizer path locally or name on huggingface",
-    )
     parser.add_argument(
         "--max-tokens",
         type=int,
-        default=512,
+        default=256,
         help="max number of tokens to generate",
     )
     parser.add_argument(
@@ -34,41 +28,41 @@ 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",
-    )
+    parser.add_argument("--cpu", action="store_true", help="Run demo with CPU only")
     args = parser.parse_args()
     return args
 
 
 def main(args):
     print(args)
+
+    if args.cpu:
+        device_map = "cpu"
+    else:
+        device_map = "auto"
+
     model = AutoModelForCausalLM.from_pretrained(
-        args.model, device_map="auto", torch_dtype="auto", trust_remote_code=True
-    )
-    tokenizer = AutoTokenizer.from_pretrained(
-        args.tokenizer or args.model, trust_remote_code=True
+        args.model, device_map=device_map, torch_dtype="auto"
     )
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
     inputs = tokenizer(
         args.prompt,
         return_tensors="pt",
-    )
+    ).to(model.device)
+
     streamer = TextStreamer(tokenizer) if args.streaming else None
     outputs = model.generate(
-        inputs.input_ids.cuda(),
+        **inputs,
         max_new_tokens=args.max_tokens,
         streamer=streamer,
-        eos_token_id=tokenizer.convert_tokens_to_ids(args.eos_token),
-        do_sample=True,
-        repetition_penalty=1.3,
-        no_repeat_ngram_size=5,
-        temperature=0.7,
-        top_k=40,
-        top_p=0.8,
+        # do_sample=True,
+        # repetition_penalty=1.3,
+        # no_repeat_ngram_size=5,
+        # temperature=0.7,
+        # top_k=40,
+        # top_p=0.8,
     )
+
     if streamer is None:
         print(tokenizer.decode(outputs[0], skip_special_tokens=True))