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")


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)

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)

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)

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)

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']}")