feat: Build class for summarizer pipeline

util/summarizer.py

@@ -1,74 +1,71 @@
-from transformers import BigBirdPegasusForConditionalGeneration, AutoTokenizer
+import gradio as gr
+from util import textproc
 
-# by default encoder-attention is `block_sparse` with num_random_blocks=3, block_size=64
-
-# TODO: add pre-trained summarizer models
-# Placeholder text for testing input
-test_text = """
-        Lorem Ipsum is simply dummy text of the printing and typesetting industry. Lorem Ipsum has been the industry's standard dummy text ever since the 1500s, when an unknown printer took a galley of type and scrambled it to make a type specimen book. It has survived not only five centuries, but also the leap into electronic typesetting, remaining essentially unchanged. It was popularised in the 1960s with the release of Letraset sheets containing Lorem Ipsum passages, and more recently with desktop publishing software like Aldus PageMaker including versions of Lorem Ipsum.
-"""
 summary_options = ["Abstract", "Background", "Claims"]
-
-def get_word_index(s, idx):
-    words = re.findall(r'\s*\S+\s*', s)
-    return sum(map(len, words[:idx])) + len(words[idx]) - len(words[idx].lstrip())
-
-
+model_names = ["huggingface/google/bigbird-pegasus-large-bigpatent",
+               "huggingface/cnicu/t5-small-booksum",
+               "huggingface/sshleifer/distilbart-cnn-6-6",
+               "huggingface/google/pegasus-xsum"]
+
+def init_models():
+    model = dict()
+    for name in model_names:
+        model[name] = gr.Interface.load(name)
+    return model
+
+    
 class PatentSummarizer():
-    def __init__(self, base_model_name="google/bigbird-pegasus-large-bigpatent"):
-        # Possible to tweak other summaries with different models in the future
-        self.model = dict()
-        self.tokenizer = dict()
-
-        self.base_model = BigBirdPegasusForConditionalGeneration.from_pretrained(base_model_name)
-        self.base_tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-
+    def __init__(self, model_collection):
+        self.model = model_collection
         self.max_word_input = 1000
 
 
-    def pipeline(patent_information, summaries_generated, min_char_abs, min_char_bg, min_char_claims):
-        # TODO: add checker if valid patent info, return None if invalid
-        # TODO: add scraper to get document
+    def pipeline(self, patent_information, summaries_generated, abstract_model, \
+                 background_model, claims_model, collate_claims, word_limit):
 
-        # TODO: add parser to get the following info from the base document:
-        abstract, background, claims = None, None, None
+        parsed_info = textproc.retrieve_parsed_doc(patent_information, 
+                                                   summaries_generated)
+        if parsed_info is None:
+            return ["[ERROR] Invalid Patent Information.", None, None]
 
+        abstract, background, claims = parsed_info
         summaries = list()
-        if "Abstract" in summaries_generated:
-            abstract_summary = summarizer.generate_abs_summary(abstract, min_char_abs)
-            summaries.append(abstract_summary)
-        else:
-            summaries.append(None)
-
-        if "Background" in summaries_generated:
-            background_summary = summarizer.generate_bg_summary(background, min_char_bg)
-            summaries.append(background_summary)
-        else:
-            summaries.append(None)
-
-        if "Claims" in summaries_generated:
-            claims_summary = summarizer.generate_claims_summary(claims, min_char_claims)
-            summaries.append(claims_summary)
-        else:
-            summaries.append(None)
-
-        return summaries
-
-
-    def generate_abs_summary(abstract, min_char_abs):
-        return "Abstract" + test_text
-
-
-    def generate_bg_summary(background, min_char_bg):
-        stop_idx = get_word_index(background, self.max_word_input)
-        inputs = self.base_tokenizer(background[0:stop_idx], 
-                                return_tensors='pt')
-        prediction = self.base_model.generate(**inputs)
-        bg_summary = self.base_tokenizer.batch_decode(prediction)
-        bg_summary = textproc.clean_text(bg_summary[0])
-
-        return bg_summary
-
 
-    def generate_claims_summary(claims, min_char_claims):
-        return "Claims" + test_text
+        try:
+            if "Abstract" in summaries_generated and abstract is not None:
+                abstract = abstract[0: textproc.get_word_index(abstract, word_limit)]
+
+                abstract_summary = self.model[abstract_model](abstract)
+                abstract_summary = textproc.post_process(abstract_summary)
+                summaries.append(abstract_summary)
+            else:
+                summaries.append(None)
+
+            if "Background" in summaries_generated and background is not None:
+                background = background[0: textproc.get_word_index(background, word_limit)]
+
+                background_summary = self.model[background_model](background)
+                background_summary = textproc.post_process(background_summary)
+                summaries.append(background_summary)
+            else:
+                summaries.append(None)
+
+            if "Claims" in summaries_generated and claims is not None:
+                if collate_claims:
+                    claims = ' '.join(claims)
+                    print(len(claims))
+                    claims = claims[0: textproc.get_word_index(claims, word_limit)]
+                    print(len(claims))
+                    claims_summary = self.model[claims_model](claims)
+                else:
+                    claims_summary = ''
+                    for claim in claims:
+                        claims_summary += self.model[claims_model](claim)
+                claims_summary = textproc.post_process(claims_summary)
+                summaries.append(claims_summary)
+            else:
+                summaries.append(None)
+
+            return summaries
+        except Exception as e:
+            return [f'[ERROR] {e}'] + [None]*(len(summaries_generated) - 1)