Update app.py

app.py

@@ -38,7 +38,7 @@ def infer(input_ids, max_length, temperature, top_k, top_p):
     return output_sequences
 
 
-
+@st.cache_data
 def chunking(book_text):
     sentences = sent_tokenize(book_text)
     segments = [] 
@@ -96,7 +96,7 @@ st.title("Book Summarization 📚")
 st.write("The almighty king of text generation, GPT-2 comes in four available sizes, only three of which have been publicly made available. Feared for its fake news generation capabilities, it currently stands as the most syntactically coherent model. A direct successor to the original GPT, it reinforces the already established pre-training/fine-tuning killer duo. From the paper: Language Models are Unsupervised Multitask Learners by Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei and Ilya Sutskever.")
 
 #book_index = st.sidebar.slider("Select Book Example", value = 0,min_value = 0, max_value=4)
-sent = st.text_area("Text", _book, height = 1550000)
+sent = st.text_area("Text", _book, height = 100000)
 max_length = st.sidebar.slider("Max Length", value = 512,min_value = 10, max_value=1024)
 temperature = st.sidebar.slider("Temperature", value = 1.0, min_value = 0.0, max_value=1.0, step=0.05)
 top_k = st.sidebar.slider("Top-k", min_value = 0, max_value=5, value = 0)
@@ -133,7 +133,8 @@ def generate_output(test_samples):
 chunked_segments = chunking(test_book[0]['book'])
 
 
-for segment in range(len(chunked_segments)):
+for i in range(len(chunked_segments)):
         
-    summaries = generate_output(segment)
-    st.write(summaries[-1])
+    summaries = generate_output(chunked_segments[i])
+    st.write(f'A summary of Segment {i}.')
+    st.success(summaries[-1])