Update Chat with PViLQA

app.py

@@ -2,26 +2,34 @@ import streamlit as st
 import torch
 from transformers import pipeline
 
-st.set_page_config(page_title="Vietnamese Legal Question Answering", page_icon="🧊", layout="centered", initial_sidebar_state="collapsed")
+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)
 
+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!')
-    st.session_state.model = pipeline("question-answering", model='Truong-Phuc/ViL-MRC-large', device=device)
+    st.session_state.model = pipeline("question-answering", model='./models/vi-mrc-large/archive/model', device=device)
 
 def get_answer(context, question):
     return st.session_state.model(context=context, question=question, max_answer_len=512)
 
-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)
+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'])
 
-context = st.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=400)
-question = st.text_area(label='Câu hỏi liên quan đến nội dung văn bản pháp luật Việt Nam ở trên:', placeholder='Vui lòng nhập câu hỏi liên quan đến nội dung văn bản pháp luật Việt Nam ở trên:...', height=100)
+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.session_state.messages.append({'role': 'user', 'content': prompt})
+    
+    respond = get_answer(context=context, question=prompt)['answer']
 
-col_1, col_2, col_3, col_4, col_5 = st.columns([1, 1, 1.3, 1, 1])
-btn_answer = col_3.button(label='Giải đáp thắc mắc', use_container_width=True)
+    with st.chat_message(name='assitant'):
+        st.markdown(respond)
 
-if btn_answer:
-    with st.spinner("Vui lòng chờ..."):
-        answer = get_answer(context=context, question=question)
-        st.success(f"Câu trả lời: {answer['answer']}")
+    st.session_state.messages.append({'role': 'assitant', 'content': respond})