Spaces:
Running
Running
| """Text chunking strategies for the RAG engine. | |
| Splitting text into smaller chunks is essential in RAG because: | |
| 1. Embedding models have context length limits (typically 256-512 tokens). | |
| 2. Small chunks provide high-precision retrieval matching. | |
| 3. Overlap between chunks ensures ideas aren't cut in half at boundaries. | |
| """ | |
| import re | |
| class TextChunk: | |
| """Represents a chunk of text derived from a source document.""" | |
| def __init__(self, chunk_id, doc_id, text, title=None, chunk_index=0, parent_text=None): | |
| self.chunk_id = chunk_id | |
| self.doc_id = doc_id | |
| self.text = text.strip() | |
| self.title = title or "" | |
| self.chunk_index = chunk_index | |
| # Parent text is the broader surrounding context (optional parent-child retrieval) | |
| self.parent_text = parent_text or self.text | |
| def to_dict(self): | |
| return { | |
| "chunk_id": self.chunk_id, | |
| "doc_id": self.doc_id, | |
| "text": self.text, | |
| "title": self.title, | |
| "chunk_index": self.chunk_index, | |
| "parent_text": self.parent_text, | |
| } | |
| def from_dict(cls, data): | |
| return cls( | |
| chunk_id=data.get("chunk_id"), | |
| doc_id=data.get("doc_id"), | |
| text=data.get("text", ""), | |
| title=data.get("title", ""), | |
| chunk_index=data.get("chunk_index", 0), | |
| parent_text=data.get("parent_text"), | |
| ) | |
| def __repr__(self): | |
| return f"<TextChunk id={self.chunk_id} doc={self.doc_id} len={len(self.text)}>" | |
| def split_text_into_chunks(text, chunk_size=400, chunk_overlap=50): | |
| """Split a string of text into overlapping word chunks. | |
| Args: | |
| text: The string content to split. | |
| chunk_size: Target number of words per chunk. | |
| chunk_overlap: Number of overlapping words between consecutive chunks. | |
| Returns: | |
| List of string chunks. | |
| """ | |
| text = (text or "").strip() | |
| if not text: | |
| return [] | |
| words = text.split() | |
| if len(words) <= chunk_size: | |
| return [text] | |
| chunks = [] | |
| step = max(1, chunk_size - chunk_overlap) | |
| for start in range(0, len(words), step): | |
| end = start + chunk_size | |
| chunk_words = words[start:end] | |
| chunks.append(" ".join(chunk_words)) | |
| if end >= len(words): | |
| break | |
| return chunks | |
| def chunk_document(document, chunk_size=400, chunk_overlap=50): | |
| """Chunk a Document into TextChunk objects. | |
| Args: | |
| document: A Document instance. | |
| chunk_size: Words per chunk. | |
| chunk_overlap: Overlap between consecutive chunks. | |
| Returns: | |
| List of TextChunk instances. | |
| """ | |
| text_chunks = split_text_into_chunks( | |
| document.text, | |
| chunk_size=chunk_size, | |
| chunk_overlap=chunk_overlap, | |
| ) | |
| chunks = [] | |
| for idx, piece in enumerate(text_chunks): | |
| cid = f"{document.doc_id}_chunk_{idx}" | |
| chunk = TextChunk( | |
| chunk_id=cid, | |
| doc_id=document.doc_id, | |
| text=piece, | |
| title=document.title, | |
| chunk_index=idx, | |
| parent_text=document.text, # Keep full document text as parent context | |
| ) | |
| chunks.append(chunk) | |
| return chunks | |
| def chunk_documents(documents, chunk_size=400, chunk_overlap=50): | |
| """Chunk multiple Document objects.""" | |
| all_chunks = [] | |
| for doc in documents: | |
| all_chunks.extend(chunk_document(doc, chunk_size, chunk_overlap)) | |
| return all_chunks | |