chat

app.py

@@ -1,22 +1,54 @@
 import streamlit as st
+from langchain.memory import ConversationBufferMemory
+from langchain.chains import ConversationalRetrievalChain
+from langchain_community.llms import HuggingFaceHub
 
+def get_conversation(model):
+    
+    memory = ConversationBufferMemory(memory_key="messages", return_messages=True)
+
+    conversation_chain = ConversationalRetrievalChain.from_llm(
+        llm=model,
+        # retriever=vectorstore.as_retriever(),
+        memory=memory)
+    
+    return conversation_chain
+
+def get_response(conversation_chain, query):
+    # get the response
+    response = conversation_chain.invoke(query)
+    response = response["result"]
+    answer = response.split('\nHelpful Answer: ')[1]
+    return answer
 
 def main():
     st.title("Health Chatbot")
     
-    if 'messages' not in st.session_state:
+    print("Loading LLM from HuggingFace")
+    with st.spinner('Loading LLM from HuggingFace...'):
+        llm = HuggingFaceHub(repo_id="HuggingFaceH4/zephyr-7b-beta", model_kwargs={"temperature":0.7, "max_new_tokens":512, "top_p":0.95, "top_k":50})
+    
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+        
+    if st.button("Clear Chat"):
         st.session_state.messages = []
     
-
-    user_input = st.chat_input("You", key="user_input")
-    if user_input:
-        st.session_state.messages.append({"sender": "user", "message": user_input})
-        # display all the messages already in the session state
-        for message in st.session_state.messages:
-            if message["sender"] == "user":
-                st.write("You: ", message["message"])
-            else:
-                st.write("Bot: ", message["message"])
-
+    for message in st.session_state.messages:
+        if message["role"] == "user":
+            st.chat_message("user").markdown(message["content"])
+        else:
+            st.chat_message("bot").markdown(message["content"])
+    
+    conversation_chain = get_conversation(llm)
+    
+    user_prompt = st.chat_input("ask a question", key="user")
+    if user_prompt:
+        st.chat_message("user").markdown(user_prompt)
+        st.session_state.messages.append({"role": "user", "content": user_prompt})
+        response = get_response(conversation_chain, user_prompt)
+        st.chat_message("bot").markdown(response)
+        st.session_state.messages.append({"role": "bot", "content": response})
+        
 if __name__ == "__main__":
     main()