Update app.py

app.py

@@ -2,13 +2,22 @@ from transformers import AutoModel, AutoTokenizer
 import gradio as gr
 import mdtex2html
 
-tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
-model = AutoModel.from_pretrained("MOSS550V/divination", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b-int4", trust_remote_code=True)
+model = AutoModel.from_pretrained("THUDM/chatglm-6b-int4", trust_remote_code=True)
 model = model.quantize(4)
 model = model.half().cuda()
 model.transformer.prefix_encoder.float()
 model = model.eval()
 
+CHECKPOINT_PATH = "MOSS550V/divination"
+
+prefix_state_dict = torch.load(os.path.join(CHECKPOINT_PATH, "pytorch_model.bin"))
+new_prefix_state_dict = {}
+for k, v in prefix_state_dict.items():
+    if k.startswith("transformer.prefix_encoder."):
+        new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v
+model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)
+
 """Override Chatbot.postprocess"""