import streamlit as st
import torch
from transformers import pipeline

st.set_page_config(page_title="Vietnamese Legal Question Answering", page_icon="🧊", layout="wide", initial_sidebar_state="collapsed")

@st.cache_data
def load_model(model_path):
    device = 0 if torch.cuda.is_available() else -1
    question_answerer = pipeline("question-answering", model=model_path, device=device)
    return question_answerer

def get_answer(model, context, question):
    return model(context=context, question=question, max_answer_len=512)

if 'model' not in st.session_state:
    st.session_state.model = load_model(model_path='./model')

st.markdown("<h1 style='text-align: center;'>Vietnamese Legal Question Answering</h1>", unsafe_allow_html=True)

context = st.text_area(label='Vietnamese Legal Documents/context:', placeholder='Enter your Vietnamese legal document here...', height=300)
question = st.text_area(label='Question about this Vietnamese Legal Documents:', placeholder='Enter your question about this Vietnamese Legal Documents here...', height=100)

btn_answer = st.button(label='Answer')

if btn_answer:
    answer = get_answer(model=st.session_state.model, context=context, question=question)
    st.success(f"{answer['answer']}")