from textwrap import wrap
from transformers import pipeline
import streamlit as st

st.markdown('# Terms & conditions abstractive summarization model :pencil:')
st.write('This app summarizes the provided terms & conditions.')
st.write('Information about the model :point_right: https://huggingface.co/ml6team/distilbart-tos-summarizer-tosdr')

st.markdown("""
To use this:
- Copy terms & conditions and hit 'Summarize':point_down:""")


@st.cache(allow_output_mutation=True,
          suppress_st_warning=True,
          show_spinner=False)
def load_model():
    with st.spinner('Please wait for the model to load...'):
        terms_and_conditions_pipeline = pipeline(
            task='summarization',
            model='ml6team/distilbart-tos-summarizer-tosdr',
            tokenizer='ml6team/distilbart-tos-summarizer-tosdr'
        )
    return terms_and_conditions_pipeline


tc_pipeline = load_model()

if 'text' not in st.session_state:
    st.session_state['text'] = ""

st.header("Input")
form = st.form(key='terms-and-conditions')
placeholder = form.empty()
placeholder.empty()
tc_text = placeholder.text_area(
    value=st.session_state.text,
    label='Terms & conditions text:',
    key='tc_text',
    height=240

)
submit_button = form.form_submit_button(label='Summarize')

st.header("Output")

if submit_button:
    base_text = st.session_state.tc_text
    output_text = " ".join([result['summary_text'] for result in tc_pipeline(wrap(base_text, 2048))])
    st.markdown('#####')
    st.text_area(
        value=output_text,
        label="Summary",
        height=240
    )