Update app.py

app.py

@@ -43,11 +43,7 @@ model_name = 'KhantKyaw/Chat_GPT-2'
 tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 model = GPT2LMHeadModel.from_pretrained(model_name)
 
-
-
-
-
-st.title("Chat with GPT-2")
+st.title("Echo Bot")
 
 # Initialize chat history
 if "messages" not in st.session_state:
@@ -55,21 +51,22 @@ if "messages" not in st.session_state:
 
 # Display chat messages from history on app rerun
 for message in st.session_state.messages:
-    with st.container():
-        st.markdown(f"**{message['role'].capitalize()}**: {message['content']}")
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
 
 # React to user input
-prompt = st.text_input("What is up?", key="chat_input")
-if prompt:
-    with st.container():
-        st.markdown(f"**User**: {prompt}")
-        st.session_state.messages.append({"role": "user", "content": prompt})
-    
-
+if prompt := st.chat_input("What is up?"):
+    # Display user message in chat message container
+    st.chat_message("user").markdown(prompt)
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # Decode the generated tokens and remove the eos token
     response = generate_response(prompt)
+    # Display assistant response in chat message container
+    with st.chat_message("assistant"):
+        st.markdown(response)
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": response})
+
+
 
-    with st.container():
-        st.markdown(f"**GPT-2**: {response}")
-        st.session_state.messages.append({"role": "assistant", "content": response})