Update app.py

app.py

@@ -1,22 +1,26 @@
-from transformers import AutoModel, AutoTokenizer
+import os, sys
+
 import gradio as gr
 import mdtex2html
 
-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()
+import torch
+import transformers
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoTokenizer,
+    AutoTokenizer,
+    DataCollatorForSeq2Seq,
+    HfArgumentParser,
+    Seq2SeqTrainingArguments,
+    set_seed,
+)
+
+from arguments import ModelArguments, DataTrainingArguments
 
-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)
+model = None
+tokenizer = None
 
 """Override Chatbot.postprocess"""
 
@@ -86,7 +90,7 @@ def reset_state():
 
 
 with gr.Blocks() as demo:
-    gr.HTML("""<h1 align="center">预测</h1>""")
+    gr.HTML("""<h1 align="center">ChatGLM</h1>""")
 
     chatbot = gr.Chatbot()
     with gr.Row():
@@ -98,9 +102,9 @@ 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=64, step=1.0, label="Maximum length", interactive=True)
+            max_length = gr.Slider(0, 4096, value=2048, 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.45, step=0.01, label="Temperature", interactive=True)
+            temperature = gr.Slider(0, 1, value=0.95, step=0.01, label="Temperature", interactive=True)
 
     history = gr.State([])
 
@@ -110,4 +114,55 @@ with gr.Blocks() as demo:
 
     emptyBtn.click(reset_state, outputs=[chatbot, history], show_progress=True)
 
-demo.queue().launch(share=False, inbrowser=True)
+
+
+def main():
+    global model, tokenizer
+
+    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(
+        model_args.model_name_or_path, trust_remote_code=True)
+    config = AutoConfig.from_pretrained(
+        model_args.model_name_or_path, trust_remote_code=True)
+
+    config.pre_seq_len = model_args.pre_seq_len
+    config.prefix_projection = model_args.prefix_projection
+
+    ptuning_checkpoint = "MOSS550V/divination"
+
+    if ptuning_checkpoint is not None:
+        print(f"Loading prefix_encoder weight from {ptuning_checkpoint}")
+        model = AutoModel.from_pretrained(model_args.model_name_or_path, 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(model_args.model_name_or_path, 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)
+
+    if model_args.pre_seq_len is not None:
+        # P-tuning v2
+        model = model.half().cuda()
+        model.transformer.prefix_encoder.float().cuda()
+    
+    model = model.eval()
+    demo.queue().launch(share=False, inbrowser=True)
+
+
+
+if __name__ == "__main__":
+    main()