data/pdf_loader.py

"""
PDF Document Ingestion, Clean Text Extraction, and Semantic Chunking.
"""
import os
import re
from typing import List, Dict, Any, Optional
from pypdf import PdfReader


class PDFDocument:
    """Represents an ingested PDF with page-level text and extracted chunks."""
    def __init__(self, filepath: str, pages: List[Dict[str, Any]], full_text: str):
        self.filepath = filepath
        self.filename = os.path.basename(filepath)
        self.pages = pages  # List of {"page_num": int, "text": str}
        self.full_text = full_text

    def __len__(self) -> int:
        return len(self.pages)


class PDFLoader:
    """Loads and cleans PDF files for model training and context retrieval."""

    @staticmethod
    def clean_text(text: str) -> str:
        """Normalize whitespace, remove weird control characters, and clean formatting."""
        # Replace non-breaking spaces and irregular whitespace
        text = re.sub(r"[\r\f\v]", "\n", text)
        # Fix hyphenated words broken across line breaks (e.g., trans-\nformer -> transformer)
        text = re.sub(r"(\w+)-\n(\w+)", r"\1\2", text)
        # Collapse multiple empty lines to max two
        text = re.sub(r"\n{3,}", "\n\n", text)
        # Normalize excessive horizontal whitespace
        text = re.sub(r"[ \t]+", " ", text)
        return text.strip()

    @classmethod
    def load_pdf(cls, filepath: str) -> PDFDocument:
        """Extract text from all pages of a PDF file."""
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"PDF file not found at: {filepath}")

        reader = PdfReader(filepath)
        pages = []
        full_text_parts = []

        for i, page in enumerate(reader.pages):
            raw_text = page.extract_text() or ""
            cleaned = cls.clean_text(raw_text)
            if cleaned:
                pages.append({
                    "page_num": i + 1,
                    "text": cleaned
                })
                full_text_parts.append(f"[Page {i+1}]\n{cleaned}")

        full_text = "\n\n".join(full_text_parts)
        return PDFDocument(filepath=filepath, pages=pages, full_text=full_text)

    @classmethod
    def chunk_document(
        cls,
        doc: PDFDocument,
        chunk_size_words: int = 150,
        chunk_overlap_words: int = 30
    ) -> List[Dict[str, Any]]:
        """
        Split document pages into overlapping semantic chunks for context retrieval and training.
        """
        chunks = []
        chunk_id = 0

        for page_data in doc.pages:
            page_num = page_data["page_num"]
            words = page_data["text"].split()

            if len(words) <= chunk_size_words:
                chunks.append({
                    "chunk_id": chunk_id,
                    "page_num": page_num,
                    "text": page_data["text"]
                })
                chunk_id += 1
            else:
                start = 0
                while start < len(words):
                    end = min(start + chunk_size_words, len(words))
                    chunk_text = " ".join(words[start:end])
                    chunks.append({
                        "chunk_id": chunk_id,
                        "page_num": page_num,
                        "text": chunk_text
                    })
                    chunk_id += 1
                    if end == len(words):
                        break
                    start += (chunk_size_words - chunk_overlap_words)

        return chunks