Included support for chunk summarize

summarize.py

@@ -1,16 +1,29 @@
 from datetime import datetime
+import multiprocessing
 from transformers import BartTokenizer, TFBartForConditionalGeneration, pipeline
 from Utils import fetch_article_text, count_tokens
 import re
 from nltk.tokenize import sent_tokenize
 import nltk
+import threading
 
 tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-cnn')
 model = TFBartForConditionalGeneration.from_pretrained('facebook/bart-large-cnn')
+max_length = model.config.max_position_embeddings
 
-def bart_summarize(text: str):
+summaries = []
+
+def generate_summary(text: str):
+    encoded_input = tokenizer.encode(text, max_length=max_length, return_tensors='tf')
+        
+    # generate summary for the input chunk
+    summary_ids = model.generate(encoded_input, max_length=300, num_beams=4, early_stopping=True)
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    
+    # add the summary to the list of summaries
+    summaries.append(summary)
 
-    max_length = model.config.max_position_embeddings
+def bart_summarize(text: str):
 
     try:
         sentences = sent_tokenize(text)
@@ -36,19 +49,17 @@ def bart_summarize(text: str):
         input_chunks.append(temp_sentences)
 
     # summarize each input chunk separately
-    summaries = []
+    results = []
     print(datetime.now().strftime("%H:%M:%S"))
     for chunk in input_chunks:
-        # encode the input chunk
+        result_t = multiprocessing.Process(target=generate_summary, args=(chunk,))
+        results.append(result_t)
+    
+    for result in results:
+        result.start()
 
-        encoded_input = tokenizer.encode(chunk, max_length=max_length, return_tensors='tf')
-        
-        # generate summary for the input chunk
-        summary_ids = model.generate(encoded_input, max_length=300, num_beams=4, early_stopping=True)
-        summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-        
-        # add the summary to the list of summaries
-        summaries.append(summary)
+    for result in results:
+        result.join()
         
     # # combine the summaries to get the final summary for the entire input
     final_summary = " ".join(summaries)