Update app.py

app.py

@@ -21,13 +21,14 @@ with st.spinner('Wait for it... the model is loading'):
     model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Create a text input for the user
-input_text = st.text_input("Enter your message:")
+# Create a chat input for the user
+input_text = st.chat_input("Enter your message:")
 
 # Check if the user has entered a message
 if input_text:
     # Add the user's message to the conversation history
     session_state.conversation_history.append(input_text)
+    st.chat_message("user", input_text)
 
     # Create conversation history string
     history_string = "\n".join(session_state.conversation_history)
@@ -43,10 +44,4 @@ if input_text:
 
     # Add the model's response to the conversation history
     session_state.conversation_history.append(response)
-
-    # Display the conversation history
-    st.write("Conversation History:")
-    for i in range(0, len(session_state.conversation_history), 2):
-        st.write("User: " + session_state.conversation_history[i])
-        if i+1 < len(session_state.conversation_history):
-            st.write("Assistant: " + session_state.conversation_history[i+1])
+    st.chat_message("assistant", response)