Update README

ZhangErling · Apr 12, 2023 · 1a368af · 1a368af
1 parent 3736c1a
commit 1a368af
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/ptuning/README.md b/ptuning/README.md
@@ -141,7 +141,7 @@ from transformers import AutoConfig, AutoModel, AutoTokenizer
 # Load model and tokenizer of ChatGLM-6B
 config = AutoConfig.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True, pre_seq_len=128)
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
-model = AutoModel.from_pretrained("THUDM/chatglm-6b", config=config, trust_remote_code=True).half().cuda()
+model = AutoModel.from_pretrained("THUDM/chatglm-6b", config=config, trust_remote_code=True)
 
 # Load PrefixEncoder
 prefix_state_dict = torch.load(os.path.join(CHECKPOINT_PATH, "pytorch_model.bin"))
@@ -150,6 +150,10 @@ for k, v in prefix_state_dict.items():
     new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v
 model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)
 
+print(f"Quantized to 4 bit")
+model = model.quantize(4)
+model = model.half().cuda()
+model.transformer.prefix_encoder.float()
 model = model.eval()
 
 response, history = model.chat(tokenizer, "你好", history=[])