Update app.py

app.py

@@ -102,7 +102,7 @@ with gr.Blocks() as demo:
                 submitBtn = gr.Button("Submit", variant="primary")
         with gr.Column(scale=1):
             emptyBtn = gr.Button("Clear History")
-            max_length = gr.Slider(0, 4096, value=2048, step=1.0, label="Maximum length", interactive=True)
+            max_length = gr.Slider(0, 4096, value=64, step=1.0, label="Maximum length", interactive=True)
             top_p = gr.Slider(0, 1, value=0.7, step=0.01, label="Top P", interactive=True)
             temperature = gr.Slider(0, 1, value=0.95, step=0.01, label="Temperature", interactive=True)
 
@@ -123,9 +123,9 @@ def main():
         ModelArguments))
 
     tokenizer = AutoTokenizer.from_pretrained(
-        "THUDM/chatglm-6b", trust_remote_code=True)
+        "THUDM/chatglm-6b-int4", trust_remote_code=True)
     config = AutoConfig.from_pretrained(
-        "THUDM/chatglm-6b", trust_remote_code=True)
+        "THUDM/chatglm-6b-int4", trust_remote_code=True)
 
     config.pre_seq_len = 128
     config.prefix_projection = False
@@ -134,15 +134,15 @@ def main():
 
     if ptuning_checkpoint is not None:
         print(f"Loading prefix_encoder weight from {ptuning_checkpoint}")
-        model = AutoModel.from_pretrained("THUDM/chatglm-6b", config=config, trust_remote_code=True)
-        prefix_state_dict = torch.load(os.path.join(ptuning_checkpoint, "pytorch_model.bin"), map_location='cpu')
+        model = AutoModel.from_pretrained("THUDM/chatglm-6b-int4", config=config, trust_remote_code=True)
+        prefix_state_dict = torch.load(os.path.join(ptuning_checkpoint, "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)
     else:
-        model = AutoModel.from_pretrained("THUDM/chatglm-6b", config=config, trust_remote_code=True)
+        model = AutoModel.from_pretrained("THUDM/chatglm-6b-int4", config=config, trust_remote_code=True)
 
     # model = model.quantize(4)