Update generate.py

generate.py

@@ -13,17 +13,18 @@ from utils import StreamPeftGenerationMixin,StreamLlamaForCausalLM
 from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
 
 parser = argparse.ArgumentParser()
-parser.add_argument("--model_path", type=str, default="/model/13B_hf")
-parser.add_argument("--lora_path", type=str, default="checkpoint-3000")
-parser.add_argument("--use_typewriter", type=int, default=1)
-parser.add_argument("--use_local", type=int, default=1)
+TOT_CUDA="0,1" #Upgrade bitsandbytes to the latest version to enable balanced loading of multiple GPUs, for example: pip install bitsandbytes==0.39.0
+BASE_MODEL="ziqingyang/chinese-llama-2-13b"
+LORA_PATH="teachyourselfcoding/llama-2-13b-22sep"
+USE_LOCAL=1 # 1: use local model, 0: use huggingface model
+TYPE_WRITER=1 # whether output streamly
+
 args = parser.parse_args()
 print(args)
 tokenizer = LlamaTokenizer.from_pretrained(args.model_path)
 
 LOAD_8BIT = True
-BASE_MODEL = args.model_path
-LORA_WEIGHTS = args.lora_path
+
 
 
 # fix the path for local checkpoint