update app.py

app.py

@@ -33,15 +33,14 @@ pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 prompt = st.text_input("Please ask your question here:")
 
-# Button to submit question
-if st.button("Submit"):
-    if prompt.strip():
+
+
+if prompt := st.chat_input("Please ask your question here:"):
         # Display user message in chat message container
-        st.markdown("<b style='color:Red;'>User:</b> {}".format(prompt), unsafe_allow_html=True)
-        
+        st.chat_message("user").markdown("<b style='color:Red;'>"+prompt+"</b>", unsafe_allow_html=True)
+        # Add user message to chat history
+        st.session_state.query = prompt
+        st.session_state.messages.append({"role": "user", "content": prompt})
         # Generate response
         result = pipe(prompt)#, max_length=100, num_return_sequences=1)
-        
-        # Display bot response in chat message container
-        st.markdown("<b style='color:Green;'>Bot:</b> {}".format(result[0]['generated_text']), unsafe_allow_html=True)
-
+