Spaces:

umarbalak
/

QueryMind

Sleeping

QueryMind / data_extraction.py

version 1

4ff3463 over 1 year ago

8.76 kB

import json
import re
from pathlib import Path

class DataExtractor:
    def __init__(self):
        """
        Initialize DataExtractor with robust logging and error handling
        """
        pass

    def extract_text_from_pdf(self, pdf_path):
        """
        Robust PDF text extraction with enhanced error handling
        
        Args:
            pdf_path (str): Path to PDF file
        
        Returns:
            str: Extracted text or empty string
        """
        try:
            import pymupdf
            doc = pymupdf.open(pdf_path)
            
            # Extract text with page-level error handling
            full_text = []
            for page_num, page in enumerate(doc, 1):
                try:
                    page_text = page.get_text()
                    full_text.append(page_text)
                except Exception as page_error:
                    print(f"Text extraction failed for page {page_num} in {pdf_path}: {page_error}")
            
            return ''.join(full_text)
        
        except Exception as e:
            print(f"PDF Extraction failed for {pdf_path}: {e}")
            return ""
        
    def extract_text_from_doc(self, doc_path):
        """
        Robust DOCX text extraction with comprehensive error handling
        
        Args:
            doc_path (str): Path to DOCX file
        
        Returns:
            str: Extracted text or empty string
        """
        try:
            import docx2txt
            return docx2txt.process(doc_path)
        except Exception as e:
            print(f"DOCX Extraction failed for {doc_path}: {e}")
            return ""
        
    def extract_text_from_txt(self, txt_path):
        """
        Robust TXT file text extraction
        
        Args:
            txt_path (str): Path to TXT file
        
        Returns:
            str: Extracted text or empty string
        """
        try:
            with open(txt_path, "r", encoding="utf-8") as f:
                return f.read()
        except UnicodeDecodeError:
            # Try alternative encodings
            for encoding in ['latin-1', 'iso-8859-1', 'cp1252']:
                try:
                    with open(txt_path, "r", encoding=encoding) as f:
                        return f.read()
                except:
                    continue
        except Exception as e:
            print(f"TXT Extraction failed for {txt_path}: {e}")
            return ""
        
    def clean_text(self, text):
        """
        Robust cleaner for government FAQ chatbot.
        - Removes boilerplate (page numbers, headers/footers)
        - Preserves layout for LLM tokenization
        - Normalizes structure and character encoding
        """

        # Normalize line endings
        text = text.replace('\r\n', '\n').replace('\r', '\n')

        # Remove page headers/footers (common boilerplate)
        text = re.sub(r'\n?Page\s*\d+(\s*of\s*\d+)?\n?', '\n', text, flags=re.IGNORECASE)
        text = re.sub(r'\n?Copyright.*?\d{4}.*?\n?', '\n', text, flags=re.IGNORECASE)
        text = re.sub(r'\n?All rights reserved.*?\n?', '\n', text, flags=re.IGNORECASE)

        # Remove decorative dividers
        text = re.sub(r'^\s*[-=_*]{3,}\s*$', '', text, flags=re.MULTILINE)

        # Strip non-printable characters but keep structure
        text = re.sub(r'[^\x09\x0A\x0D\x20-\x7E]', '', text)

        # Normalize whitespace
        text = re.sub(r'[ \t]+', ' ', text)             # inline
        text = re.sub(r'[ \t]+\n', '\n', text)          # line ends
        text = re.sub(r'\n{3,}', '\n\n', text)          # excessive newlines

        return text.strip()

    
    def extract_text_from_files(self, doc_paths):
        """
        Extract and process text from multiple document types
        
        Returns:
            list: Extracted document data with metadata
        """
        doc_paths = [
            str(Path(path).resolve()) 
            for path in doc_paths 
            if Path(path).exists()
        ]
        data = []

        for path in doc_paths:
            try:
                ext = Path(path).suffix.lower()
                
                # Text extraction based on file type
                if ext == '.pdf':
                    text = self.extract_text_from_pdf(path)
                elif ext == '.docx':
                    text = self.extract_text_from_doc(path)
                elif ext == '.txt':
                    text = self.extract_text_from_txt(path)
                else:
                    print(f"Unsupported format: {path}")
                    continue
                
                # Skip empty documents
                if not text.strip():
                    print(f"No text extracted from {path}")
                    continue

                # Clean and structure extracted text
                cleaned_text = self.clean_text(text)
                
                # Add document metadata
                doc_data = {
                    "title": Path(path).name,
                    "path": path,
                    "text": cleaned_text,
                    "text_length": len(cleaned_text)
                }
                
                data.append(doc_data)

            except Exception as e:
                print(f"Unexpected error processing {path}: {e}")

        return data
    
    def pdf_to_markdown(self, pdf_path, markdown_path):
        """
        Convert PDF to Markdown using docling
        
        Args:
            pdf_path (Path): Path to PDF file
            markdown_path (Path): Output path for markdown file
            
        Returns:
            str: Markdown content
        """
        try:
            from docling.document_converter import DocumentConverter

            # Define the source PDF file
            source = pdf_path
            converter = DocumentConverter()

            # Convert the PDF to Markdown
            result = converter.convert(source)
            markdown = result.document.export_to_markdown()

            # Write the Markdown output to a file
            with open(markdown_path, "w", encoding="utf-8") as file:
                file.write(markdown)

            return markdown
        except ImportError:
            print("Docling is not installed. Please install it with: pip install docling")
            return ""
        except Exception as e:
            print(f"Error converting PDF to Markdown: {e}")
            return ""

    def load_markdown_file(self, markdown_path):
        """
        Load markdown file using langchain_docling
        
        Args:
            markdown_path (Path): Path to markdown file
            
        Returns:
            list: List of documents
        """
        try:
            from langchain_docling import DoclingLoader
            from langchain_docling.loader import ExportType

            loader = DoclingLoader(
                file_path=markdown_path,
                export_type=ExportType.MARKDOWN
            )
            data = loader.load()
            return data
        except ImportError:
            print("langchain_docling is not installed. Please install it with: pip install langchain_docling")
            return []
        except Exception as e:
            print(f"Error loading markdown file: {e}")
            return []
    
    def spit_markdown(self, markdown):
        """
        Split markdown using MarkdownHeaderTextSplitter
        
        Args:
            markdown (list): List of markdown documents
            
        Returns:
            list: List of split documents
        """
        try:
            from langchain_text_splitters import MarkdownHeaderTextSplitter

            splitter = MarkdownHeaderTextSplitter(
                headers_to_split_on=[
                    ("#", "Header_1"),
                    ("##", "Header_2"),
                    ("###", "Header_3"),
                ],
            )

            splits = []
            for doc in markdown:
                for split in splitter.split_text(doc.page_content):
                    # Include metadata and page_content from the original document
                    split.metadata.update({
                        "Header_1": split.metadata.get("Header_1", ""),
                        "Header_2": split.metadata.get("Header_2", ""),
                        "Header_3": split.metadata.get("Header_3", "")
                    })
                    splits.append(split)
            
            return splits
        except ImportError:
            print("langchain_text_splitters is not installed or has compatibility issues.")
            return []
        except Exception as e:
            print(f"Error splitting markdown: {e}")
            return []

	import json
	import re
	from pathlib import Path

	class DataExtractor:
	def __init__(self):
	"""
	Initialize DataExtractor with robust logging and error handling
	"""
	pass

	def extract_text_from_pdf(self, pdf_path):
	"""
	Robust PDF text extraction with enhanced error handling

	Args:
	pdf_path (str): Path to PDF file

	Returns:
	str: Extracted text or empty string
	"""
	try:
	import pymupdf
	doc = pymupdf.open(pdf_path)

	# Extract text with page-level error handling
	full_text = []
	for page_num, page in enumerate(doc, 1):
	try:
	page_text = page.get_text()
	full_text.append(page_text)
	except Exception as page_error:
	print(f"Text extraction failed for page {page_num} in {pdf_path}: {page_error}")

	return ''.join(full_text)

	except Exception as e:
	print(f"PDF Extraction failed for {pdf_path}: {e}")
	return ""

	def extract_text_from_doc(self, doc_path):
	"""
	Robust DOCX text extraction with comprehensive error handling

	Args:
	doc_path (str): Path to DOCX file

	Returns:
	str: Extracted text or empty string
	"""
	try:
	import docx2txt
	return docx2txt.process(doc_path)
	except Exception as e:
	print(f"DOCX Extraction failed for {doc_path}: {e}")
	return ""

	def extract_text_from_txt(self, txt_path):
	"""
	Robust TXT file text extraction

	Args:
	txt_path (str): Path to TXT file

	Returns:
	str: Extracted text or empty string
	"""
	try:
	with open(txt_path, "r", encoding="utf-8") as f:
	return f.read()
	except UnicodeDecodeError:
	# Try alternative encodings
	for encoding in ['latin-1', 'iso-8859-1', 'cp1252']:
	try:
	with open(txt_path, "r", encoding=encoding) as f:
	return f.read()
	except:
	continue
	except Exception as e:
	print(f"TXT Extraction failed for {txt_path}: {e}")
	return ""

	def clean_text(self, text):
	"""
	Robust cleaner for government FAQ chatbot.
	- Removes boilerplate (page numbers, headers/footers)
	- Preserves layout for LLM tokenization
	- Normalizes structure and character encoding
	"""

	# Normalize line endings
	text = text.replace('\r\n', '\n').replace('\r', '\n')

	# Remove page headers/footers (common boilerplate)
	text = re.sub(r'\n?Page\s\d+(\sof\s*\d+)?\n?', '\n', text, flags=re.IGNORECASE)
	text = re.sub(r'\n?Copyright.?\d{4}.?\n?', '\n', text, flags=re.IGNORECASE)
	text = re.sub(r'\n?All rights reserved.*?\n?', '\n', text, flags=re.IGNORECASE)

	# Remove decorative dividers
	text = re.sub(r'^\s[-=_]{3,}\s*$', '', text, flags=re.MULTILINE)

	# Strip non-printable characters but keep structure
	text = re.sub(r'[^\x09\x0A\x0D\x20-\x7E]', '', text)

	# Normalize whitespace
	text = re.sub(r'[ \t]+', ' ', text) # inline
	text = re.sub(r'[ \t]+\n', '\n', text) # line ends
	text = re.sub(r'\n{3,}', '\n\n', text) # excessive newlines

	return text.strip()


	def extract_text_from_files(self, doc_paths):
	"""
	Extract and process text from multiple document types

	Returns:
	list: Extracted document data with metadata
	"""
	doc_paths = [
	str(Path(path).resolve())
	for path in doc_paths
	if Path(path).exists()
	]
	data = []

	for path in doc_paths:
	try:
	ext = Path(path).suffix.lower()

	# Text extraction based on file type
	if ext == '.pdf':
	text = self.extract_text_from_pdf(path)
	elif ext == '.docx':
	text = self.extract_text_from_doc(path)
	elif ext == '.txt':
	text = self.extract_text_from_txt(path)
	else:
	print(f"Unsupported format: {path}")
	continue

	# Skip empty documents
	if not text.strip():
	print(f"No text extracted from {path}")
	continue

	# Clean and structure extracted text
	cleaned_text = self.clean_text(text)

	# Add document metadata
	doc_data = {
	"title": Path(path).name,
	"path": path,
	"text": cleaned_text,
	"text_length": len(cleaned_text)
	}

	data.append(doc_data)

	except Exception as e:
	print(f"Unexpected error processing {path}: {e}")

	return data

	def pdf_to_markdown(self, pdf_path, markdown_path):
	"""
	Convert PDF to Markdown using docling

	Args:
	pdf_path (Path): Path to PDF file
	markdown_path (Path): Output path for markdown file

	Returns:
	str: Markdown content
	"""
	try:
	from docling.document_converter import DocumentConverter

	# Define the source PDF file
	source = pdf_path
	converter = DocumentConverter()

	# Convert the PDF to Markdown
	result = converter.convert(source)
	markdown = result.document.export_to_markdown()

	# Write the Markdown output to a file
	with open(markdown_path, "w", encoding="utf-8") as file:
	file.write(markdown)

	return markdown
	except ImportError:
	print("Docling is not installed. Please install it with: pip install docling")
	return ""
	except Exception as e:
	print(f"Error converting PDF to Markdown: {e}")
	return ""

	def load_markdown_file(self, markdown_path):
	"""
	Load markdown file using langchain_docling

	Args:
	markdown_path (Path): Path to markdown file

	Returns:
	list: List of documents
	"""
	try:
	from langchain_docling import DoclingLoader
	from langchain_docling.loader import ExportType

	loader = DoclingLoader(
	file_path=markdown_path,
	export_type=ExportType.MARKDOWN
	)
	data = loader.load()
	return data
	except ImportError:
	print("langchain_docling is not installed. Please install it with: pip install langchain_docling")
	return []
	except Exception as e:
	print(f"Error loading markdown file: {e}")
	return []

	def spit_markdown(self, markdown):
	"""
	Split markdown using MarkdownHeaderTextSplitter

	Args:
	markdown (list): List of markdown documents

	Returns:
	list: List of split documents
	"""
	try:
	from langchain_text_splitters import MarkdownHeaderTextSplitter

	splitter = MarkdownHeaderTextSplitter(
	headers_to_split_on=[
	("#", "Header_1"),
	("##", "Header_2"),
	("###", "Header_3"),
	],
	)

	splits = []
	for doc in markdown:
	for split in splitter.split_text(doc.page_content):
	# Include metadata and page_content from the original document
	split.metadata.update({
	"Header_1": split.metadata.get("Header_1", ""),
	"Header_2": split.metadata.get("Header_2", ""),
	"Header_3": split.metadata.get("Header_3", "")
	})
	splits.append(split)

	return splits
	except ImportError:
	print("langchain_text_splitters is not installed or has compatibility issues.")
	return []
	except Exception as e:
	print(f"Error splitting markdown: {e}")
	return []