71 lines
2.4 KiB
Python
71 lines
2.4 KiB
Python
"""
|
||
Text splitters for chunking documents.
|
||
"""
|
||
|
||
from enum import Enum
|
||
from typing import List, Optional
|
||
|
||
from langchain_core.documents import Document
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain_experimental.text_splitter import SemanticChunker
|
||
|
||
|
||
class SplitterType(str, Enum):
|
||
RECURSIVE = "recursive"
|
||
SEMANTIC = "semantic"
|
||
PARENT_CHILD = "parent_child"
|
||
|
||
|
||
def get_splitter(splitter_type: SplitterType, **kwargs):
|
||
"""Factory function to create a text splitter."""
|
||
if splitter_type == SplitterType.RECURSIVE:
|
||
chunk_size = kwargs.get("chunk_size", 500)
|
||
chunk_overlap = kwargs.get("chunk_overlap", 50)
|
||
return RecursiveCharacterTextSplitter(
|
||
chunk_size=chunk_size,
|
||
chunk_overlap=chunk_overlap,
|
||
separators=["\n\n", "\n", "。", "！", "？", " ", ""],
|
||
)
|
||
elif splitter_type == SplitterType.SEMANTIC:
|
||
# Requires embeddings for semantic splitting
|
||
embeddings = kwargs.get("embeddings")
|
||
if embeddings is None:
|
||
raise ValueError("Semantic splitter requires 'embeddings' parameter")
|
||
return SemanticChunker(embeddings=embeddings)
|
||
else:
|
||
raise ValueError(f"Unsupported splitter type: {splitter_type}")
|
||
|
||
|
||
class ParentChildSplitter:
|
||
"""
|
||
Splits documents into parent (large) and child (small) chunks.
|
||
Child chunks are indexed for retrieval, parent chunks are stored for context.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
parent_chunk_size: int = 1000,
|
||
child_chunk_size: int = 200,
|
||
parent_chunk_overlap: int = 100,
|
||
child_chunk_overlap: int = 20,
|
||
):
|
||
self.parent_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=parent_chunk_size,
|
||
chunk_overlap=parent_chunk_overlap,
|
||
)
|
||
self.child_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=child_chunk_size,
|
||
chunk_overlap=child_chunk_overlap,
|
||
)
|
||
|
||
def split_documents(self, documents: List[Document]) -> tuple[List[Document], List[Document]]:
|
||
"""
|
||
Returns:
|
||
(parent_chunks, child_chunks)
|
||
"""
|
||
parent_chunks = self.parent_splitter.split_documents(documents)
|
||
child_chunks = self.child_splitter.split_documents(documents)
|
||
|
||
# Link child chunks to parent IDs (optional metadata)
|
||
# In a real implementation, you'd map each child to a parent chunk ID.
|
||
return parent_chunks, child_chunks |