Update app.py

app.py

@@ -4,14 +4,30 @@ import mdtex2html
 import torch
 import os
 
-#CHECKPOINT_PATH=f'output_lh/checkpoint-600'
+
+CHECKPOINT_PATH=f'output_lh_v2/checkpoint-700'
 tokenizer = AutoTokenizer.from_pretrained("chatglm3-6b", trust_remote_code=True)
+config = AutoConfig.from_pretrained("chatglm3-6b", trust_remote_code=True, pre_seq_len=128)
+model = AutoModel.from_pretrained("chatglm3-6b", config=config, trust_remote_code=True)
+prefix_state_dict = torch.load(os.path.join(CHECKPOINT_PATH, "pytorch_model.bin"), map_location=torch.device('cpu'))
+
+#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)
 
-model = AutoModel.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True, device='cuda')
+#model = model.half().cuda()
 
+model = model.float()
+
+model.transformer.prefix_encoder.float()
 model = model.eval()
 
 
+
 """Override Chatbot.postprocess"""