Update README.md

README.md

@@ -41,19 +41,16 @@ from transformers import AutoTokenizer
 from faster_chat_glm import GLM6B, FasterChatGLM
 
 
-tokenizer = AutoTokenizer.from_pretrained(chatglm6b_dir, trust_remote_code=True)
-
-BATCH_SIZE = 8
-MAX_OUT_LEN = 50
-
-# prepare input
-input_str = ["为什么我们需要对深度学习模型加速? ", ] *
+MAX_OUT_LEN = 100
+tokenizer = AutoTokenizer.from_pretrained('./models', trust_remote_code=True)
+input_str = ["为什么我们需要对深度学习模型加速？", ]
 inputs = tokenizer(input_str, return_tensors="pt", padding=True)
 input_ids = inputs.input_ids.to('cuda:0')
 
 
+plan_path = './models/glm6b-bs8.ftm'
 # kernel for chat model.
-kernel = GLM6B(plan_path="./models/glm6b-bs{BATCH_SIZE}.ftm",
+kernel = GLM6B(plan_path=plan_path,
                batch_size=1,
                num_beams=1,
                use_cache=True,
@@ -62,7 +59,8 @@ kernel = GLM6B(plan_path="./models/glm6b-bs{BATCH_SIZE}.ftm",
                decoder_layers=28,
                vocab_size=150528,
                max_seq_len=MAX_OUT_LEN)
-chat = FasterChatGLM(model_dir=chatglm6b_dir, kernel=kernel).half().cuda()
+
+chat = FasterChatGLM(model_dir="./models", kernel=kernel).half().cuda()
 
 # generate
 sample_output = chat.generate(inputs=input_ids, max_length=MAX_OUT_LEN)