Update app.py

app.py

@@ -32,7 +32,45 @@ def infer(input_ids, max_length, temperature, top_k, top_p):
     
     return output_sequences
 
-default_value = test_book[0]['book']
+
+def chunking(book_text):
+    segments = [] 
+    #sentences, token_lens
+    current_segment = ""    
+    total_token_lens = 0
+    for i in range(len(sentences)):
+        
+        if total_token_lens < 512:
+            total_token_lens += token_lens[i]
+            current_segment += (sentences[i] + " ")
+            
+        elif total_token_lens > 768:
+            segments.append(current_segment)
+            current_segment = sentences[i]
+            total_token_lens = token_lens[i]
+    
+        else:
+            #make next_pseudo_segment
+            next_pseudo_segment = ""
+            next_token_len = 0
+            for t in range(30):
+                if (i+t < len(sentences)) and (next_token_len + token_lens[i+t] < 512):
+                    next_token_len += token_lens[i+t] 
+                    next_pseudo_segment += sentences[i+t] 
+                
+            embs = model.encode([current_segment, next_pseudo_segment, sentences[i]]) # current, next, sent
+            if cos_similarity(embs[1],embs[2]) > cos_similarity(embs[0],embs[2]):
+                segments.append(current_segment)
+                current_segment = sentences[i]
+                total_token_lens = token_lens[i]
+            else: 
+                total_token_lens += token_lens[i]
+                current_segment += (sentences[i] + " ")
+
+    return segments
+
+
+chunked_segments = chunking(test_book[0]['book'])
 
 '''
 '''
@@ -41,40 +79,47 @@ default_value = test_book[0]['book']
 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.")
 
-sent = st.text_area("Text", default_value, height = 550)
+book_index = st.sidebar.slider("Select Book Example", value = 0,min_value = 0, max_value=4)
+
+_book = test_book[book_index]['book']
+chunked_segments = chunking(_book)
+
+sent = st.text_area("Text", _book, height = 550)
 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)
 top_p = st.sidebar.slider("Top-p", min_value = 0.0, max_value=1.0, step = 0.05, value = 0.92)
 
-encoded_prompt = tokenizer.encode(sent, add_special_tokens=False, return_tensors="pt")
-if encoded_prompt.size()[-1] == 0:
-    input_ids = None
-else:
-    input_ids = encoded_prompt
-
-
-output_sequences = infer(input_ids, max_length, temperature, top_k, top_p)
-
-
-
-for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
-    print(f"=== GENERATED SEQUENCE {generated_sequence_idx + 1} ===")
-    generated_sequences = generated_sequence.tolist()
-
-    # Decode text
-    text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
-
-    # Remove all text after the stop token
-    #text = text[: text.find(args.stop_token) if args.stop_token else None]
-
-    # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
-    total_sequence = (
-        sent + text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)) :]
-    )
 
-    generated_sequences.append(total_sequence)
-    print(total_sequence)
-
-
-st.write(generated_sequences[-1])
+for segment in range(len(chunked_segments)):
+        
+    encoded_prompt = tokenizer.encode(segment, add_special_tokens=False, return_tensors="pt")
+    if encoded_prompt.size()[-1] == 0:
+        input_ids = None
+    else:
+        input_ids = encoded_prompt
+    
+    
+    output_sequences = infer(input_ids, max_length, temperature, top_k, top_p)
+    
+    
+    for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
+        print(f"=== GENERATED SEQUENCE {generated_sequence_idx + 1} ===")
+        generated_sequences = generated_sequence.tolist()
+    
+        # Decode text
+        text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
+    
+        # Remove all text after the stop token
+        #text = text[: text.find(args.stop_token) if args.stop_token else None]
+    
+        # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
+        total_sequence = (
+            sent + text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)) :]
+        )
+    
+        generated_sequences.append(total_sequence)
+        print(total_sequence)
+    
+    
+    st.write(generated_sequences[-1])