import streamlit as st

mapper = {
    "wikilingua": "64b53c9e04b2bfaeb2f0b38cb7712bcbd4755c3f",
    "vietnews": "d8d516ad2c112b429155c3e6077182ae5fe5b33d"
}

if "model" not in st.session_state:
    from src.model import Model
    st.session_state.task="wikilingua"
    model = Model(revision=mapper["wikilingua"])
    st.session_state.model = model
    
    
st.set_page_config(page_title="ViT5 Reproduce", layout="wide")
hide_menu_style = """
<style>
footer {visibility: hidden;}
</style>
"""
st.markdown(hide_menu_style, unsafe_allow_html= True)

with st.sidebar:
    task = st.selectbox(label="Task", options=["wikilingua", "vietnews"])
    if task != st.session_state.task:
        from src.model import Model
        st.session_state.task=task
        model = Model(revision=mapper[task])
    
left, middle, right = st.columns([4,1,4])

left_container = left.container(border=True)
left_container.write(f"**Input**")
left_container.divider()
text = left_container.text_area(label="", height=512, label_visibility="hidden", max_chars=4096*5)

summary_button = middle.button("Summary ➩", type="primary", use_container_width=True)

right_container = right.container(border=True)
right_container.markdown(f"**Output**")
right_container.divider()

if summary_button:
    
    output = st.session_state.model.inference(text=text)
                                          
                                          
    st.session_state["output"] = output

if "output" in st.session_state:
    right_container.text_area(label="", value=st.session_state["output"], height=512, label_visibility="hidden")