Tidy Up Code

.gitignore

@@ -0,0 +1,4 @@
+src/test.py
+test
+test.py
+test.sh

src/app.py

@@ -1,72 +1,17 @@
-from src.text_extractor import TextExtractor
-from tqdm import tqdm
-from transformers import PegasusForConditionalGeneration, PegasusTokenizer
-from transformers import pipeline
-from mdutils.mdutils import MdUtils
-from pathlib import Path
-
+from src.summarizer import Summarizer
 import gradio as gr
-import fitz
-import torch
-import copy
-import os
-
-FILENAME = ""
-
-preprocess = TextExtractor()
-model_name = "sshleifer/distill-pegasus-cnn-16-4"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-tokenizer = PegasusTokenizer.from_pretrained(model_name)
-model = PegasusForConditionalGeneration.from_pretrained(model_name).to(device)
-
-def summarize(slides):
-    generated_slides = copy.deepcopy(slides)
-    for page, contents in tqdm(generated_slides.items()):
-        for idx, (tag, content) in enumerate(contents):
-            if tag.startswith('p'): 
-                try:
-                    input = tokenizer(content, truncation=True, padding="longest", return_tensors="pt").to(device)
-                    tensor = model.generate(**input)
-                    summary = tokenizer.batch_decode(tensor, skip_special_tokens=True)[0]
-                    contents[idx] = (tag, summary)
-                except Exception as e:
-                    print(f"Summarization Fails, Error: {e}")
-    return generated_slides
-
-def convert2markdown(generated_slides):
-    mdFile = MdUtils(file_name=FILENAME, title=f'{FILENAME} Presentation')
-    for k, v in generated_slides.items():
-        mdFile.new_line('---\n')
-        for section in v:
-            tag = section[0]
-            content = section[1]
-            if tag.startswith('h'):
-                mdFile.new_header(level=int(tag[1]), title=content)
-            if tag == 'p':
-                contents = content.split('<n>')
-                for content in contents:
-                    mdFile.new_line(f"{content}\n")
-    mdFile.create_md_file()
-    return f"{FILENAME}.md"
 
 def inference(document):
-    global FILENAME
-    doc = fitz.open(document)
-    FILENAME = document.name.split('/')[-1].split('.')[0]
-    font_counts, styles = preprocess.get_font_info(doc, granularity=False)
-    size_tag = preprocess.get_font_tags(font_counts, styles)
-    texts = preprocess.assign_tags(doc, size_tag)
-    slides = preprocess.get_slides(texts)
-    generated_slides = summarize(slides)
-    markdown_name = convert2markdown(generated_slides)
-    print(f"Markdown File Name: {markdown_name}")
-    return markdown_name
-
+    summarizer = Summarizer("sshleifer/distill-pegasus-cnn-16-4")
+    slide_content = summarizer.extract_text(document)
+    summarized_slides = summarizer(slide_content)
+    markdown = summarizer.convert2markdown(summarized_slides)
+    print(f"Markdown File Name: {markdown.file_name}")
+    return markdown.file_name
 
 with gr.Blocks() as demo:
     inp = gr.File(file_types=['pdf'])
     out = gr.File(label="Markdown File")
-    # out = gr.Textbox(label="Markdown Content")
     inference_btn = gr.Button("Summarized PDF")
     inference_btn.click(fn=inference, inputs=inp, outputs=out, show_progress=True, api_name="summarize")
     

src/summarizer.py

@@ -0,0 +1,62 @@
+from typing import Dict, List, Tuple, Optional
+from tqdm import tqdm
+from transformers import PegasusForConditionalGeneration, PegasusTokenizer
+from src.text_extractor import TextExtractor
+from mdutils.mdutils import MdUtils
+
+import torch
+import fitz
+import copy
+
+class Summarizer():
+    def __init__(self, model_name: str):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.tokenizer = PegasusTokenizer.from_pretrained(model_name)
+        self.model = PegasusForConditionalGeneration.from_pretrained(model_name).to(self.device)
+        self.preprocess = TextExtractor()
+    
+    def extract_text(self, document: object) -> Dict[str, List[Tuple[str, str]]]:
+        doc = fitz.open(document)
+        self.filename = doc.name.split('/')[-1].split('.')[0]
+        font_counts, styles = self.preprocess.get_font_info(doc, granularity=False)
+        size_tag = self.preprocess.get_font_tags(font_counts, styles)
+        texts = self.preprocess.assign_tags(doc, size_tag)
+        slide_content = self.preprocess.get_slides(texts)
+        return slide_content
+
+    def __call__(self, slides: Dict[str, List[Tuple[str, str]]]) -> Dict[str, List[Tuple[str, str]]]:
+            summarized_slides = copy.deepcopy(slides)
+            for page, contents in tqdm(summarized_slides.items()):
+                for idx, (tag, content) in enumerate(contents):
+                    if tag.startswith('p'): 
+                        try:
+                            input = self.tokenizer(content, truncation=True, padding="longest", return_tensors="pt").to(self.device)
+                            tensor = self.model.generate(**input)
+                            summary = self.tokenizer.batch_decode(tensor, skip_special_tokens=True)[0]
+                            contents[idx] = (tag, summary)
+                        except Exception as e:
+                            print(f"Summarization Fails, Error: {e}")
+                            
+            return summarized_slides
+    
+    def convert2markdown(self, summarized_slides: Dict[str, List[Tuple[str, str]]], target_path: Optional[str]=None) -> str:
+        filename = self.filename
+        if target_path:
+            filename = target_path
+        mdFile = MdUtils(file_name=filename, title=f'{self.filename} Presentation')
+        for k, v in summarized_slides.items():
+            mdFile.new_line('---\n')
+            for section in v:
+                tag = section[0]
+                content = section[1]
+                if tag.startswith('h'):
+                    mdFile.new_header(level=int(tag[1]), title=content)
+                if tag == 'p':
+                    contents = content.split('<n>')
+                    for content in contents:
+                        mdFile.new_line(f"{content}\n")
+        markdown = mdFile.create_md_file()
+        return markdown
+
+        
+        

src/text_extractor.py

@@ -2,7 +2,6 @@ from operator import itemgetter
 from collections import OrderedDict
 from typing import Dict, List, Iterator, Union, Tuple
 
-
 import re
 
 class TextExtractor: