Update app.py

app.py

@@ -3,7 +3,7 @@
 from transformers import AutoModel, AutoTokenizer
 import streamlit as st
 from streamlit_chat import message
-
+from fastllm_pytools import llm
 
 st.set_page_config(
     page_title="ChatGLM2-6b 演示",
@@ -16,7 +16,9 @@ st.set_page_config(
 def get_model():
     tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True)
     model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True).float()
-    model = model.eval()
+
+    model = llm.from_hf(model, tokenizer, dtype = "int4")
+    #model = model.eval()
     return tokenizer, model