Spaces:
Sleeping
Sleeping
File size: 5,934 Bytes
40e5eae | 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 | """
Text Splitting Module
Splits documents into chunks for embedding and retrieval
"""
from pathlib import Path
from typing import List, Dict
import logging
from langchain_text_splitters import RecursiveCharacterTextSplitter
from config import (
PROCESSED_DOCS_DIR,
TEXT_SPLITTER_CONFIG,
)
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class DocumentChunker:
"""
Handles document chunking with configurable parameters
"""
def __init__(
self,
chunk_size: int = TEXT_SPLITTER_CONFIG["chunk_size"],
chunk_overlap: int = TEXT_SPLITTER_CONFIG["chunk_overlap"],
separators: List[str] = TEXT_SPLITTER_CONFIG["separators"],
keep_separator: bool = TEXT_SPLITTER_CONFIG["keep_separator"],
):
"""
Initialize the document chunker
Args:
chunk_size: Maximum size of each chunk in characters
chunk_overlap: Number of characters to overlap between chunks
separators: List of separator strings to split on
keep_separator: Whether to keep the separator in the chunks
"""
self.splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
separators=separators,
keep_separator=keep_separator,
)
logger.info(
f"Initialized text splitter: chunk_size={chunk_size}, "
f"chunk_overlap={chunk_overlap}"
)
def split_text(self, text: str) -> List[str]:
"""
Split text into chunks
Args:
text: Input text to split
Returns:
List of text chunks
"""
try:
chunks = self.splitter.split_text(text)
logger.debug(f"Split text into {len(chunks)} chunks")
return chunks
except Exception as e:
logger.error(f"Error splitting text: {e}")
raise
def split_document(self, file_path: Path) -> List[Dict[str, str]]:
"""
Split a markdown document into chunks with metadata
Args:
file_path: Path to the markdown file
Returns:
List of dictionaries containing chunk text and metadata
"""
try:
# Read the file
with open(file_path, 'r', encoding='utf-8') as f:
text = f.read()
# Split into chunks
chunks = self.split_text(text)
# Add metadata to each chunk
chunk_data = []
for idx, chunk in enumerate(chunks):
chunk_data.append({
"text": chunk,
"source": file_path.name,
"chunk_id": idx,
"total_chunks": len(chunks),
})
logger.info(f"Split {file_path.name} into {len(chunks)} chunks")
return chunk_data
except Exception as e:
logger.error(f"Error splitting document {file_path}: {e}")
raise
def process_all_documents() -> List[Dict[str, str]]:
"""
Process all markdown documents in the processed_docs directory
Returns:
List of all chunks with metadata from all documents
"""
chunker = DocumentChunker()
all_chunks = []
if not PROCESSED_DOCS_DIR.exists():
logger.error(f"Processed documents directory not found: {PROCESSED_DOCS_DIR}")
return all_chunks
# Process all markdown files
markdown_files = list(PROCESSED_DOCS_DIR.glob("*.md"))
if not markdown_files:
logger.warning("No markdown files found in processed_docs directory")
return all_chunks
logger.info(f"Processing {len(markdown_files)} documents...")
for file_path in markdown_files:
try:
chunks = chunker.split_document(file_path)
all_chunks.extend(chunks)
logger.info(f"Added {len(chunks)} chunks from {file_path.name}")
except Exception as e:
logger.error(f"Failed to process {file_path.name}: {e}")
continue
logger.info(f"Total chunks processed: {len(all_chunks)}")
return all_chunks
def get_chunk_statistics(chunks: List[Dict[str, str]]) -> Dict:
"""
Calculate statistics about the chunks
Args:
chunks: List of chunk dictionaries
Returns:
Dictionary with statistics
"""
if not chunks:
return {
"total_chunks": 0,
"total_characters": 0,
"avg_chunk_size": 0,
"min_chunk_size": 0,
"max_chunk_size": 0,
"sources": [],
}
chunk_sizes = [len(chunk["text"]) for chunk in chunks]
sources = list(set(chunk["source"] for chunk in chunks))
return {
"total_chunks": len(chunks),
"total_characters": sum(chunk_sizes),
"avg_chunk_size": sum(chunk_sizes) / len(chunk_sizes),
"min_chunk_size": min(chunk_sizes),
"max_chunk_size": max(chunk_sizes),
"sources": sources,
"num_sources": len(sources),
}
if __name__ == "__main__":
# Test the text splitting
logger.info("Testing text splitting...")
chunks = process_all_documents()
if chunks:
stats = get_chunk_statistics(chunks)
logger.info(f"Chunk statistics: {stats}")
# Display first chunk as example
if chunks:
logger.info("\nExample chunk:")
logger.info(f"Source: {chunks[0]['source']}")
logger.info(f"Chunk ID: {chunks[0]['chunk_id']}")
logger.info(f"Text preview: {chunks[0]['text'][:200]}...")
else:
logger.warning("No chunks created. Make sure documents are converted first.")
|