Update app.py

app.py

@@ -10,6 +10,7 @@ import gradio as gr
 import mdtex2html
 
 model_name = "THUDM/chatglm2-6b"
+model_name = "THUDM/chatglm2-6b-int4"
 
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
@@ -24,9 +25,9 @@ has_cuda = torch.cuda.is_available()
 # has_cuda = False  # force cpu
 
 if has_cuda:
-    model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4",trust_remote_code=True).cuda()  # 3.92
+    model = AutoModel.from_pretrained(model_name,trust_remote_code=True).cuda()  # 3.92
 else:
-    model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4",trust_remote_code=True).float()
+    model = AutoModel.from_pretrained(model_name,trust_remote_code=True).float()
 
 model = model.eval()
 
@@ -100,7 +101,7 @@ def reset_state():
 
 
 with gr.Blocks() as demo:
-    gr.HTML("""<h1 align="center">ChatGLM2-6B</h1>""")
+    gr.HTML("""<h1 align="center">ChatGLM2-6B-4bit</h1>""")
 
     chatbot = gr.Chatbot()
     with gr.Row():