Update app.py

app.py

@@ -3,13 +3,17 @@ import torch
 from transformers import pipeline
 
 st.set_page_config(page_title="Vietnamese Legal Question Answering", page_icon="🧊", layout="centered", initial_sidebar_state="expanded")
-st.markdown("<h1 style='text-align: center;'>Hệ thống hỏi đáp trực tuyến cho văn bản pháp luật Việt Nam</h1>", unsafe_allow_html=True)
+
+with open("./static/styles.css") as f:
+    st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
+
+st.markdown("<h1 style='text-align: center;'>ChatViLD: Hệ thống hỏi đáp trực tuyến cho văn bản pháp luật Việt Nam</h1>", unsafe_allow_html=True)
 
 context = st.sidebar.text_area(label='Nội dung văn bản pháp luật Việt Nam:', placeholder='Vui lòng nhập nội dung văn bản pháp luật Việt Nam tại đây...', height=925)
 
 device = 0 if torch.cuda.is_available() else -1
 if 'model' not in st.session_state:
-    print('Some errors occured!')
+    print('Some errors occurred!')
     st.session_state.model = pipeline("question-answering", model='./model', device=device)
 
 def get_answer(context, question):
@@ -19,17 +23,36 @@ if 'messages' not in st.session_state:
     st.session_state.messages = []
 
 for message in st.session_state.messages:
-    with st.chat_message(name=message['role']):
-        st.markdown(message['content'])
+    if message['role'] == 'assistant':
+        avatar_class = "assistant-avatar"
+        message_class = "assistant-message"
+        avatar = './app/static/ai.png'
+    else:
+        avatar_class = "user-avatar"
+        message_class = "user-message"
+        avatar = './app/static/human.png'
+    st.markdown(f"""
+    <div class="{message_class}">
+        <img src="{avatar}" class="{avatar_class}" />
+        <div class="stMarkdown">{message['content']}</div>
+    </div>
+    """, unsafe_allow_html=True)
 
 if prompt := st.chat_input(placeholder='Xin chào, tôi có thể giúp được gì cho bạn?'):
-    with st.chat_message(name='user'):
-        st.markdown(prompt)
+    st.markdown(f"""
+    <div class="user-message">
+        <img src="./app/static/human.png" class="user-avatar" />
+        <div class="stMarkdown">{prompt}</div>
+    </div>
+    """, unsafe_allow_html=True)
     st.session_state.messages.append({'role': 'user', 'content': prompt})
     
     respond = get_answer(context=context, question=prompt)['answer']
 
-    with st.chat_message(name='assitant'):
-        st.markdown(respond)
-
-    st.session_state.messages.append({'role': 'assitant', 'content': respond})
+    st.markdown(f"""
+    <div class="assistant-message">
+        <img src="./app/static/ai.png" class="assistant-avatar" />
+        <div class="stMarkdown">{respond}</div>
+    </div>
+    """, unsafe_allow_html=True)
+    st.session_state.messages.append({'role': 'assistant', 'content': respond})