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="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='./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)

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 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']

    with st.chat_message(name='assitant'):
        st.markdown(respond)

    st.session_state.messages.append({'role': 'assitant', 'content': respond})