From 1f96081b743941b85dbec5b7ae646391ccbb0e63 Mon Sep 17 00:00:00 2001 From: Junyang Lin Date: Thu, 3 Aug 2023 13:43:07 +0800 Subject: [PATCH] Update demo.py --- demo.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/demo.py b/demo.py index 2950f05..38272df 100644 --- a/demo.py +++ b/demo.py @@ -5,8 +5,7 @@ import torch import argparse -from pathlib import Path -from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig +from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.trainer_utils import set_seed @@ -68,7 +67,7 @@ def demo_qwen_chat(args): if __name__ == "__main__": parser = argparse.ArgumentParser(description="Test HF checkpoint.") - parser.add_argument("-c", "--checkpoint-path", type=Path, help="Checkpoint path") + parser.add_argument("-c", "--checkpoint-path", type=str, help="Checkpoint path") parser.add_argument("-s", "--seed", type=int, default=1234, help="Random seed") parser.add_argument("--gpu", type=int, default=0, help="gpu id") @@ -78,4 +77,4 @@ if __name__ == "__main__": if 'chat' in args.checkpoint_path.lower(): demo_qwen_chat(args) else: - demo_qwen_pretrain(args) \ No newline at end of file + demo_qwen_pretrain(args)