Update app.py

app.py

@@ -1,74 +1,54 @@
 import streamlit as st
-from queue import Queue
-from langchain import HuggingFaceHub, PromptTemplate, LLMChain
-
-# Set the title of the Streamlit app
-st.title("Falcon QA Bot")
-
-# Get the Hugging Face Hub API token from Streamlit secrets
-huggingfacehub_api_token = st.secrets["hf_token"]
-
-# Set the repository ID for the Falcon model
-repo_id = "tiiuae/falcon-7b-instruct"
-
-# Initialize the Hugging Face Hub and LLMChain
-llm = HuggingFaceHub(
-    huggingfacehub_api_token=huggingfacehub_api_token,
-    repo_id=repo_id,
-    model_kwargs={"temperature": 0.2, "max_new_tokens": 2000}
-)
-
-# Define the template for the assistant's response
-template = """
-You are an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.
-
-{question}
-"""
-
-# Create a queue to store user questions
-queue = Queue()
-
-def chat(query):
-    """
-    Generates a response to the user's question using the LLMChain model.
-
-    :param query: User's question.
-    :type query: str
-    :return: Response to the user's question.
-    :rtype: str
-    """
-    # Create a prompt template with the question variable
-    prompt = PromptTemplate(template=template, input_variables=["question"])
-
-    # Create an LLMChain instance with the prompt and the Falcon model
-    llm_chain = LLMChain(prompt=prompt, verbose=True, llm=llm)
-
-    # Generate a response to the user's question
-    result = llm_chain.predict(question=query)
-
-    return result
-
-def main():
-    """
-    Main function for the Streamlit app.
-    """
-    # Get the user's question from the input text box
-    user_question = st.text_input("What do you want to ask about", placeholder="Input your question here")
-
-    if user_question:
-        # Add the user's question to the queue
-        queue.put(user_question)
-
-        # Check if there are any waiting users
-        if not queue.empty():
-            # Get the next user's question from the queue
-            query = queue.get()
-
-            # Generate a response to the user's question
-            response = chat(query)
-
-            # Display the response to the user
-            st.write(response, unsafe_allow_html=True)
-
-if __name__ == '__main__':
-    main()
+from transformers import AutoModel, AutoTokenizer
+import mdtex2html
+from utils import load_model_on_gpus
+
+st.set_page_config(page_title="ChatGLM2-6B", page_icon=":robot:") 
+
+st.header("ChatGLM2-6B")
+
+tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
+model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
+# Load model on multiple GPUs
+#model = load_model_on_gpus("THUDM/chatglm2-6b", num_gpus=2)  
+model = model.eval()
+
+def postprocess(chat):
+    for i, (user, response) in enumerate(chat):
+        chat[i] = (mdtex2html.convert(user), mdtex2html.convert(response))
+    return chat
+
+user_input = st.text_area("Input:", height=200, placeholder="Ask me anything!")
+if user_input:
+    history = st.session_state.get('history', [])
+    
+    max_length = st.slider("Max Length:", 0, 32768, 8192, 1)  
+    top_p = st.slider("Top P:", 0.0, 1.0, 0.8, 0.01)
+    temperature = st.slider("Temperature:", 0.0, 1.0, 0.95, 0.01)
+
+    if 'past_key_values' not in st.session_state:
+        st.session_state['past_key_values'] = None
+
+    with st.spinner("Thinking..."):
+        response = model.generate(tokenizer.encode(user_input), 
+                                max_length=max_length,  
+                                top_p=top_p, 
+                                temperature=temperature,
+                                return_dict_in_generate=True,
+                                output_scores=True,
+                                return_past_key_values=True,
+                                past_key_values=st.session_state.past_key_values)
+
+    st.session_state.past_key_values = response.past_key_values
+    
+    history.append((user_input, response.sequences[0]))
+    history = postprocess(history)
+
+    for user, chatbot in history:
+        message = f"**Human:** {user}" if user else ""
+        response = f"**AI:** {chatbot}" if chatbot else ""  
+        st.markdown(message + response, unsafe_allow_html=True) 
+
+if st.button("Clear History"):
+    st.session_state['history'] = []
+    st.session_state['past_key_values'] = None