Update app.py

app.py

@@ -43,33 +43,33 @@ model_name = 'KhantKyaw/Chat_GPT-2'
 tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 model = GPT2LMHeadModel.from_pretrained(model_name)
 
-# Chat loop
-#print("Chatbot is ready. Type 'quit' to exit.")
-#while True:
-    #user_input = input("You: ")
-    #if user_input.lower() == "quit":
-        #break
-    #response = generate_response(user_input)
-    #print("Chatbot:", response)
 
 
 
 
-prompt = st.text_input("Say Something!", key=None, max_chars=None, disabled=False)
+st.title("Chat with GPT-2")
 
-if prompt:
+# Initialize chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display chat messages from history on app rerun
+for message in st.session_state.messages:
     with st.container():
-        # Displaying the user's input question.
-        st.markdown(prompt)
-        
-        # Generating and displaying the response.
-        response = generate_response(prompt)
-        st.markdown(generate_response(prompt))
+        st.markdown(f"**{message['role'].capitalize()}**: {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})
+    
 
 
+    # Decode the generated tokens and remove the eos token
+    response = generate_response(prompt)
 
-#prompt = st.chat_input(placeholder="Say Something!",key=None, max_chars=None, disabled=False, on_submit=None, args=None, kwargs=None)
-#if prompt:
-  # with st.chat_message(name="AI",avatar=None):
-   # st.write(generate_response(prompt))
+    with st.container():
+        st.markdown(f"**GPT-2**: {response}")
+        st.session_state.messages.append({"role": "assistant", "content": response})