Update app.py

app.py

@@ -121,26 +121,20 @@ def main():
 
     parser = HfArgumentParser((
         ModelArguments))
-    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
-        # If we pass only one argument to the script and it's the path to a json file,
-        # let's parse it to get our arguments.
-        model_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))[0]
-    else:
-        model_args = parser.parse_args_into_dataclasses()[0]
 
     tokenizer = AutoTokenizer.from_pretrained(
         "THUDM/chatglm-6b-int4", trust_remote_code=True)
     config = AutoConfig.from_pretrained(
         "MOSS550V/divination", trust_remote_code=True)
 
-    config.pre_seq_len = model_args.pre_seq_len
-    config.prefix_projection = model_args.prefix_projection
+    config.pre_seq_len = 128
+    config.prefix_projection = false
 
     ptuning_checkpoint = "MOSS550V/divination"
 
     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)
+        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():
@@ -150,14 +144,11 @@ def main():
     else:
         model = AutoModel.from_pretrained("THUDM/chatglm-6b-int4", config=config, trust_remote_code=True)
 
-    if model_args.quantization_bit is not None:
-        print(f"Quantized to {model_args.quantization_bit} bit")
-        model = model.quantize(model_args.quantization_bit)
+    model = model.quantize(4)
 
-    if model_args.pre_seq_len is not None:
-        # P-tuning v2
-        model = model.half()
-        model.transformer.prefix_encoder.float()
+    # P-tuning v2
+    model = model.half()
+    model.transformer.prefix_encoder.float()
     
     model = model.eval()
     demo.queue().launch(share=False, inbrowser=True)