Update app.py

app.py

@@ -22,6 +22,15 @@ task = "text-generation"  # Change this to your model's task
 # Load the model using the pipeline
 model_pipeline = pipeline(task, model=model,tokenizer=tokenizer)
 
+def render_chat_history(chat_history):
+    #renders chat history
+    for message in chat_history:
+        if(message["role"]!= "system"):
+            with st.chat_message(message["role"]):
+                st.write(message["content"])
+
+
+message_container = st.empty()
 
 #Application 
 with st.container():
@@ -43,22 +52,15 @@ with st.container():
 
     if 'model' not in st.session_state:
         st.session_state.model = model
-
-    #renders chat history
-    for message in st.session_state.chat_history:
-        if(message["role"]!= "system"):
-            with st.chat_message(message["role"]):
-                st.write(message["content"])
-            
     
     #Set up input text field
     input_text = st.chat_input(placeholder="Here you can chat with our hotel booking model.")
 
     if input_text:
-        with st.chat_message("user"):
+        #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
-
+        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= chat_model)
         #first_answer = chat_response.split("Human")[0] #Because of Predict it prints the whole conversation.Here we seperate the first answer only.
         tokenized_chat = tokenizer.apply_chat_template(st.session_state.chat_history, tokenize=True, add_generation_prompt=True, return_tensors="pt")
@@ -66,7 +68,10 @@ with st.container():
         outputs = model.generate(tokenized_chat, max_new_tokens=128) 
         first_answer = tokenizer.decode(outputs[0][tokenized_chat.shape[1]:],skip_special_tokens=True)
         
-        with st.chat_message("assistant"):
+        #with st.chat_message("assistant"):
             #st.write(first_answer)
-            st.session_state.chat_history.append({"role": "assistant", "content": first_answer})
-    st.markdown('</div>', unsafe_allow_html=True)
+        st.session_state.chat_history.append({"role": "assistant", "content": first_answer})    
+    st.markdown('</div>', unsafe_allow_html=True)
+
+with message_container:
+    render_chat_history(st.session_state.chat_history)