Use int4

app.py

@@ -1,20 +1,22 @@
-import torch
+import psutil
 import gradio as gr
 
 from functools import partial
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
+mem = psutil.virtual_memory()
+
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
-model = AutoModelForSeq2SeqLM.from_pretrained("THUDM/chatglm-6b", revision="int4", trust_remote_code=True).cpu().float()
+
+model = AutoModelForSeq2SeqLM.from_pretrained("THUDM/chatglm-6b", revison="int4", trust_remote_code=True).cpu().float()
 
 def chat(query, history=[]):
-    response = model.chat(tokenizer, query, history, max_length=32)
-    chat_list = [(history[i], history[i + 1]) for i in range(0, len(history) - 1, 2)]
-    return chat_list, history
+    _, history = model.chat(tokenizer, query, history, max_length=512)
+    return history, history
 
-description = "This is an unofficial chatbot application based on open source model ChatGLM-6B, running on cpu(hence max_length is limited to 32)."
+description = "This is an unofficial chatbot application based on open source model ChatGLM-6B(https://github.com/THUDM/ChatGLM-6B), running on cpu(therefore max_length is limited to 512). \nIf you want to use this chat bot in your space, 'Duplicate this space' by click the button close to 'Linked Models'. \n"
 title = "ChatGLM-6B Chatbot"
-examples = [["Hello?"]]
+examples = [["Hello?"], ["你好。"], ["介绍清华"]]
 
 chatbot_interface = gr.Interface(
     fn=chat,

requirements.txt

@@ -4,3 +4,4 @@ protobuf>=3.19.5,<3.20.1
 transformers>=4.26.1
 icetk
 cpm_kernels
+psutil