Update app.py

app.py

@@ -1,6 +1,11 @@
 import streamlit as st
 import chatbot_bedrock as demo_chat
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from langchain.schema import (
+    HumanMessage,
+    SystemMessage,
+)
+from langchain_community.chat_models.huggingface import ChatHuggingFace
 
 st.title("Hi, I am Chatbot Philio :mermaid:")
 st.write("I am your hotel booking assistant for today.")
@@ -11,9 +16,9 @@ st.write("I am your hotel booking assistant for today.")
 # base="light"
 # primaryColor="#6b4bff"
 
-tokenizer, model = demo_chat.load_model()
+model = demo_chat.load_model()
 
-model_wrapper = demo_chat.HuggingFaceModelWrapper(model, tokenizer)
+chat_model = ChatHuggingFace(llm=model)
 
 #Application 
 with st.container():
@@ -36,14 +41,14 @@ with st.container():
             st.write(message["content"])
 
     #Set up input text field
-    input_text = st.chat_input(placeholder="Here you can chat with Llamma 2 model.")
+    input_text = st.chat_input(placeholder="Here you can chat with our hotel booking model.")
 
     if input_text:
         with st.chat_message("user"):
             st.write(input_text)
             st.session_state.chat_history.append({"role" : "user", "content" : input_text}) #append message to chat history
 
-        chat_response = demo_chat.demo_chain(input_text=input_text, memory=st.session_state.memory, model= model_wrapper)
+        chat_response = demo_chat.demo_chain(input_text=input_text, memory=st.session_state.memory, model= chat_model)
         first_answer = chat_response.split("Human")[0] #Because of Predict it prints the whole conversation.Here we seperate the first answer only.
 
         with st.chat_message("assistant"):