diff --git a/open_notebook/graphs/content_processing/pdf.py b/open_notebook/graphs/content_processing/pdf.py new file mode 100644 index 0000000..e842a67 --- /dev/null +++ b/open_notebook/graphs/content_processing/pdf.py @@ -0,0 +1,150 @@ +import re +import unicodedata + +import fitz # type: ignore +from loguru import logger + +from open_notebook.graphs.content_processing.state import SourceState + +# todo: find tables - https://pymupdf.readthedocs.io/en/latest/the-basics.html#extracting-tables-from-a-page +# todo: what else can we do to make the text more readable? +# todo: try to fix encoding for some PDF that is still breaking +# def _extract_text_from_pdf(pdf_path): +# doc = fitz.open(pdf_path) +# text = "" +# logger.debug(f"Found {len(doc)} pages in PDF") +# for page in doc: +# # Use encode/decode if you need to clean up any encoding issues +# text += page.get_text().encode('utf-8').decode('utf-8') +# doc.close() +# return text + +SUPPORTED_FITZ_TYPES = [ + "application/pdf", + "application/epub+zip", +] + + +def clean_pdf_text(text): + """ + Clean text extracted from PDFs with enhanced space handling. + Preserves special characters like (, ), %, = that are valid in code/math. + + Args: + text (str): The raw text extracted from a PDF + Returns: + str: Cleaned text with minimal necessary spacing + """ + if not text: + return text + + # Step 1: Normalize Unicode characters + text = unicodedata.normalize("NFKC", text) + + # Step 2: Replace common PDF artifacts + replacements = { + # Common ligatures + "fi": "fi", + "fl": "fl", + "ff": "ff", + "ffi": "ffi", + "ffl": "ffl", + # Quotation marks and apostrophes + """: "'", """: "'", + '"': '"', + "′": "'", + "‚": ",", + "„": '"', + # Dashes and hyphens + "‒": "-", + "–": "-", + "—": "-", + "―": "-", + # Other common replacements + "…": "...", + "•": "*", + "°": " degrees ", + "¹": "1", + "²": "2", + "³": "3", + "©": "(c)", + "®": "(R)", + "™": "(TM)", + } + for old, new in replacements.items(): + text = text.replace(old, new) + + # Step 3: Clean control characters while preserving essential whitespace and special chars + text = "".join( + char + for char in text + if unicodedata.category(char)[0] != "C" + or char in "\n\t " + or char in "()%=[]{}#$@!?.,;:+-*/^<>&|~" + ) + + # Step 4: Enhanced space cleaning + text = re.sub(r"[ \t]+", " ", text) # Consolidate horizontal whitespace + text = re.sub(r" +\n", "\n", text) # Remove spaces before newlines + text = re.sub(r"\n +", "\n", text) # Remove spaces after newlines + text = re.sub(r"\n\t+", "\n", text) # Remove tabs at start of lines + text = re.sub(r"\t+\n", "\n", text) # Remove tabs at end of lines + text = re.sub(r"\t+", " ", text) # Replace tabs with single space + + # Step 5: Remove empty lines while preserving paragraph structure + text = re.sub(r"\n{3,}", "\n\n", text) # Max two consecutive newlines + text = re.sub(r"^\s+", "", text) # Remove leading whitespace + text = re.sub(r"\s+$", "", text) # Remove trailing whitespace + + # Step 6: Clean up around punctuation + text = re.sub(r"\s+([.,;:!?)])", r"\1", text) # Remove spaces before punctuation + text = re.sub(r"(\()\s+", r"\1", text) # Remove spaces after opening parenthesis + text = re.sub( + r"\s+([.,])\s+", r"\1 ", text + ) # Ensure single space after periods and commas + + # Step 7: Remove zero-width and invisible characters + text = re.sub(r"[\u200b\u200c\u200d\ufeff\u200e\u200f]", "", text) + + # Step 8: Fix hyphenation and line breaks + text = re.sub( + r"(?<=\w)-\s*\n\s*(?=\w)", "", text + ) # Remove hyphenation at line breaks + + return text.strip() + + +def _extract_text_from_pdf(pdf_path): + doc = fitz.open(pdf_path) + try: + text = "" + logger.debug(f"Found {len(doc)} pages in PDF") + for page in doc: + text += page.get_text() + normalized_text = clean_pdf_text(text) + return normalized_text + finally: + doc.close() + + +def extract_pdf(state: SourceState): + """ + Parse the text file and print its content. + """ + return_dict = {} + assert state.get("file_path"), "No file path provided" + assert state.get("identified_type") in SUPPORTED_FITZ_TYPES, "Unsupported File Type" + if ( + state.get("file_path") is not None + and state.get("identified_type") in SUPPORTED_FITZ_TYPES + ): + file_path = state.get("file_path") + try: + text = _extract_text_from_pdf(file_path) + return_dict["content"] = text + except FileNotFoundError: + raise FileNotFoundError(f"File not found at {file_path}") + except Exception as e: + raise Exception(f"An error occurred: {e}") + + return return_dict