Update app.py

app.py

@@ -6,9 +6,11 @@ from huggingface_hub import login
 import re
 import os
 
+# Load Hugging Face token
 HF_TOKEN = os.getenv("HF_TOKEN")
 login(token=HF_TOKEN)
 
+# Define models
 MODELS = {
     "athena-1": {
         "name": "🦁 Atlas-Flash",
@@ -20,6 +22,10 @@ MODELS = {
     },
 }
 
+# Profile pictures
+USER_PFP = "https://huggingface.co/front/assets/avatars.png"  # Hugging Face user avatar
+AI_PFP = "ai_pfp.png"  # Replace with the path to your AI's image or a URL
+
 class AtlasInferenceApp:
     def __init__(self):
         if "current_model" not in st.session_state:
@@ -148,16 +154,21 @@ class AtlasInferenceApp:
 
         st.markdown("*⚠️ CAUTION: Atlas is an experimental model and this is just a preview. Responses may not be expected. Please double-check sensitive information!*")
 
+        # Display chat history
         for message in st.session_state.chat_history:
-            with st.chat_message(message["role"]):
+            with st.chat_message(
+                message["role"],
+                avatar=USER_PFP if message["role"] == "user" else AI_PFP
+            ):
                 st.markdown(message["content"])
 
+        # Input box for user messages
         if prompt := st.chat_input("Message Atlas..."):
             st.session_state.chat_history.append({"role": "user", "content": prompt})
-            with st.chat_message("user"):
+            with st.chat_message("user", avatar=USER_PFP):
                 st.markdown(prompt)
 
-            with st.chat_message("assistant"):
+            with st.chat_message("assistant", avatar=AI_PFP):
                 with st.spinner("Generating response..."):
                     response = self.respond(prompt, max_tokens, temperature, top_p, top_k)
                     st.markdown(response)