File size: 3,498 Bytes
6cacf86 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """
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
|