Update to correct model pipeline

app.py

@@ -1,15 +1,14 @@
-from textwrap import wrap 
-
+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('# 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 and conditions and hit 'Summarize':point_down:""")
+- Copy terms & conditions and hit 'Summarize':point_down:""")
 
 
 @st.cache(allow_output_mutation=True,
@@ -18,7 +17,7 @@ To use this:
 def load_model():
     with st.spinner('Please wait for the model to load...'):
         terms_and_conditions_pipeline = pipeline(
-            'text2text-generation',
+            task='summarization',
             model='ml6team/distilbart-tos-summarizer-tosdr',
             tokenizer='ml6team/distilbart-tos-summarizer-tosdr'
         )
@@ -30,20 +29,28 @@ 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')
+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 and conditions text:', key='tc_text')
+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')
 
-right_area.header("Output")
+st.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:")
+    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
+    )