Update app.py

app.py

@@ -52,7 +52,7 @@ def chunking(book_text):
             #make next_pseudo_segment
             next_pseudo_segment = ""
             next_token_len = 0
-            for t in range(30):
+            for t in range(10):
                 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] 
@@ -69,8 +69,6 @@ def chunking(book_text):
     return segments
 
 
-chunked_segments = chunking(test_book[0]['book'])
-
 '''
 '''
 
@@ -83,42 +81,37 @@ book_index = st.sidebar.slider("Select Book Example", value = 0,min_value = 0, m
 _book = test_book[book_index]['book']
 chunked_segments = chunking(_book)
 
-sent = st.text_area("Text", _book, height = 550)
+sent = st.text_area("Text", _book[:512], 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)
 
+def generate_output(test_samples):
+    inputs = tokenizer(
+        test_samples,
+        padding=max_length,
+        truncation=True,
+        max_length=1024,
+        return_tensors="pt",
+    )
+    input_ids = inputs.input_ids.to(model.)
+    attention_mask = inputs.attention_mask.to(model)
+    outputs = model.generate(input_ids,
+                             max_length = 256,
+                             min_length=32,
+                             top_p = 0.92,
+                             num_beams=5,
+                             no_repeat_ngram_size=2,
+                             attention_mask=attention_mask)
+    output_str = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    return outputs, output_str
+
+
+chunked_segments = chunking(test_book[0]['book'])
+
 
 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])
+    summaries = generate_output(segment)
+    st.write(summaries[-1])