from textwrap import wrap 

from transformers import pipeline
import streamlit as st

st.markdown('# Terms and conditions abstractive summarization model :pencil:')
st.write('This app summarizes the provided terms and conditions. '
         'Terms and conditions Summarization model is based on sshleifer/distilbart-cnn-6-6')

st.markdown("""
To use this:
- Copy terms and 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(
            'text2text-generation',
            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'] = ""

left_area, right_area = st.columns(2)

left_area.header("Input")
form = left_area.form(key='terms-and-conditions')
placeholder = form.empty()
placeholder.empty()
tc_text = placeholder.text_area(value=st.session_state.text, label='Terms and conditions text:', key='tc_text')
submit_button = form.form_submit_button(label='Summarize')

right_area.header("Output")

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