Spaces:
Running on Zero
Running on Zero
File size: 8,522 Bytes
0828c2c | 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 | """Document processing module for loading and chunking documents."""
import logging
from pathlib import Path
from bs4 import BeautifulSoup
from docx import Document as DocxDocument
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from .config_loader import get_config
logger = logging.getLogger(__name__)
class DocumentProcessor:
"""Process various document formats and chunk them for RAG."""
def __init__(self):
"""Initialize document processor with configuration."""
self.config = get_config()
self.chunk_size = self.config.get("document_processing.chunk_size", 1000)
self.chunk_overlap = self.config.get("document_processing.chunk_overlap", 200)
self.supported_extensions = self.config.get(
"document_processing.supported_extensions",
[".pdf", ".docx", ".doc", ".html", ".htm", ".txt", ".md"],
)
self.text_splitter = RecursiveCharacterTextSplitter(
chunk_size=self.chunk_size,
chunk_overlap=self.chunk_overlap,
length_function=len,
separators=["\n\n", "\n", ". ", " ", ""],
)
def load_pdf(self, file_path: Path) -> list[Document]:
"""Load and extract text from PDF file.
Args:
file_path: Path to PDF file
Returns:
List of Document objects
"""
try:
loader = PyPDFLoader(str(file_path))
documents = loader.load()
logger.info(f"Loaded {len(documents)} pages from {file_path.name}")
return documents
except Exception as e:
logger.error(f"Error loading PDF {file_path}: {e}")
return []
def load_docx(self, file_path: Path) -> list[Document]:
"""Load and extract text from Word document.
Args:
file_path: Path to DOCX file
Returns:
List of Document objects
"""
try:
doc = DocxDocument(str(file_path))
text = "\n\n".join([paragraph.text for paragraph in doc.paragraphs if paragraph.text])
# Extract tables if present
if doc.tables:
for table in doc.tables:
table_text = "\n".join(
["\t".join([cell.text for cell in row.cells]) for row in table.rows]
)
text += f"\n\n{table_text}"
document = Document(
page_content=text,
metadata={"source": str(file_path), "type": "docx"},
)
logger.info(f"Loaded Word document {file_path.name}")
return [document]
except Exception as e:
logger.error(f"Error loading DOCX {file_path}: {e}")
return []
def load_html(self, file_path: Path) -> list[Document]:
"""Load and extract text from HTML file.
Args:
file_path: Path to HTML file
Returns:
List of Document objects
"""
try:
with open(file_path, encoding="utf-8") as f:
html_content = f.read()
soup = BeautifulSoup(html_content, "lxml")
# Remove script and style elements
for script in soup(["script", "style"]):
script.decompose()
# Get text
text = soup.get_text(separator="\n")
# Clean up whitespace
lines = (line.strip() for line in text.splitlines())
text = "\n".join(line for line in lines if line)
document = Document(
page_content=text,
metadata={"source": str(file_path), "type": "html"},
)
logger.info(f"Loaded HTML file {file_path.name}")
return [document]
except Exception as e:
logger.error(f"Error loading HTML {file_path}: {e}")
return []
def load_text(self, file_path: Path) -> list[Document]:
"""Load text file (txt, md, etc).
Args:
file_path: Path to text file
Returns:
List of Document objects
"""
try:
loader = TextLoader(str(file_path), encoding="utf-8")
documents = loader.load()
logger.info(f"Loaded text file {file_path.name}")
return documents
except Exception as e:
logger.error(f"Error loading text file {file_path}: {e}")
return []
def load_document(self, file_path: Path) -> list[Document]:
"""Load document based on file extension.
Args:
file_path: Path to document
Returns:
List of Document objects
"""
extension = file_path.suffix.lower()
if extension == ".pdf":
return self.load_pdf(file_path)
elif extension in [".docx", ".doc"]:
return self.load_docx(file_path)
elif extension in [".html", ".htm"]:
return self.load_html(file_path)
elif extension in [".txt", ".md"]:
return self.load_text(file_path)
else:
logger.warning(f"Unsupported file type: {extension} for {file_path.name}")
return []
def chunk_documents(self, documents: list[Document]) -> list[Document]:
"""Split documents into chunks.
Args:
documents: List of Document objects
Returns:
List of chunked Document objects
"""
if not documents:
return []
try:
chunks = self.text_splitter.split_documents(documents)
logger.info(f"Split {len(documents)} documents into {len(chunks)} chunks")
return chunks
except Exception as e:
logger.error(f"Error chunking documents: {e}")
return []
def process_directory(self, directory_path: str | Path) -> list[Document]:
"""Process all documents in a directory.
Args:
directory_path: Path to directory containing documents
Returns:
List of chunked Document objects
"""
directory = Path(directory_path)
if not directory.exists():
logger.error(f"Directory not found: {directory}")
return []
# Get main document path; optionally exclude it from vector store
main_doc_path = self.config.get("main_document.path", "")
main_doc_file = Path(main_doc_path).resolve() if main_doc_path else None
exclude_main_doc = self.config.get("main_document.exclude_from_index", True)
all_documents = []
skipped_main_doc = False
# Find all supported files
for ext in self.supported_extensions:
files = list(directory.rglob(f"*{ext}"))
for file_path in files:
# Skip main document when exclude_from_index is enabled (default)
if (
exclude_main_doc
and main_doc_file
and file_path.resolve() == main_doc_file
):
if not skipped_main_doc:
logger.info(
f"Skipping main document: {file_path.name} "
"(loaded directly, not stored in vector DB)"
)
skipped_main_doc = True
continue
logger.info(f"Processing {file_path.name}...")
docs = self.load_document(file_path)
all_documents.extend(docs)
if not all_documents:
logger.warning(f"No documents found in {directory}")
return []
logger.info(f"Loaded {len(all_documents)} documents from {directory}")
# Chunk all documents
chunked_documents = self.chunk_documents(all_documents)
return chunked_documents
def process_documents(directory_path: str | None = None) -> list[Document]:
"""Convenience function to process documents.
Args:
directory_path: Path to documents directory (uses config default if None)
Returns:
List of chunked Document objects
"""
config = get_config()
if directory_path is None:
directory_path = config.get_env("DOCUMENTS_DIR", "./data/documents")
processor = DocumentProcessor()
return processor.process_directory(directory_path)
|