zykrix
Re-added phidata as a normal folder (final fix)
c96b98a
Raw
History Blame Contribute Delete
2.41 kB
from typing import List
from phi.document.base import Document
from phi.document.chunking.strategy import ChunkingStrategy
class FixedSizeChunking(ChunkingStrategy):
"""Chunking strategy that splits text into fixed-size chunks with optional overlap"""
def __init__(self, chunk_size: int = 5000, overlap: int = 0):
# overlap must be lesser than chunk size
if overlap >= chunk_size:
raise ValueError(f"Invalid parameters: overlap ({overlap}) must be less than chunk size ({chunk_size}).")
self.chunk_size = chunk_size
self.overlap = overlap
def chunk(self, document: Document) -> List[Document]:
"""Split document into fixed-size chunks with optional overlap"""
content = self.clean_text(document.content)
content_length = len(content)
chunked_documents: List[Document] = []
chunk_number = 1
chunk_meta_data = document.meta_data
# If the document length is less than overlap, it cannot be chunked.
if len(content) <= self.overlap:
return [document]
# run the chunking only if the length of the content is greater than the overlap.
start = 0
while start + self.overlap < content_length:
end = min(start + self.chunk_size, content_length)
# Ensure we're not splitting a word in half
if end < content_length:
while end > start and content[end] not in [" ", "\n", "\r", "\t"]:
end -= 1
# If the entire chunk is a word, then just split it at chunk_size
if end == start:
end = start + self.chunk_size
chunk = content[start:end]
meta_data = chunk_meta_data.copy()
meta_data["chunk"] = chunk_number
chunk_id = None
if document.id:
chunk_id = f"{document.id}_{chunk_number}"
elif document.name:
chunk_id = f"{document.name}_{chunk_number}"
meta_data["chunk_size"] = len(chunk)
chunked_documents.append(
Document(
id=chunk_id,
name=document.name,
meta_data=meta_data,
content=chunk,
)
)
chunk_number += 1
start = end - self.overlap
return chunked_documents