Update app.py

app.py

@@ -1,16 +1,38 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+# Show title and description.
+st.title("💬 Healthcare Chatbot")
+st.write(
+    "This is a simple chatbot that uses the Llama3-Med42-8B model to generate responses. "
+    "To use this app, simply type your question in the input field below."
+)
+
+# Create a session state variable to store the chat messages. This ensures that the
+# messages persist across reruns.
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display the existing chat messages via `st.chat_message`.
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
 # Načtení modelu a tokenizeru
 model_name = "m42-health/Llama3-Med42-8B"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-st.title('Healthcare Chatbot')
+# Create a chat input field to allow the user to enter a message. This will display
+# automatically at the bottom of the page.
+if user_input := st.chat_input("What is up?"):
+
+    # Store and display the current prompt.
+    st.session_state.messages.append({"role": "user", "content": user_input})
+    with st.chat_message("user"):
+        st.markdown(user_input)
 
-# Uživatelský vstup
-user_input = st.text_input("You:", "")
-if user_input:
+    # Prepare input for the model
     messages = [
         {"role": "system", "content": (
             "You are a helpful, respectful and honest medical assistant. "
@@ -29,4 +51,7 @@ if user_input:
     output_ids = model.generate(input_ids, max_length=512, do_sample=True, temperature=0.4, top_k=150, top_p=0.75)
     response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
-    st.text_area("Bot:", response[len(input_text):])
+    # Display the model's response
+    with st.chat_message("assistant"):
+        st.markdown(response[len(input_text):])
+    st.session_state.messages.append({"role": "assistant", "content": response[len(input_text):]})