-
Notifications
You must be signed in to change notification settings - Fork 1
/
strategies.py
18 lines (16 loc) · 807 Bytes
/
strategies.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from .chunking_strategy import ChunkingStrategy
from .layout import LayoutDocumentChunking
from .page import PageDocumentChunking
from .fixed_size_overlap import FixedSizeOverlapDocumentChunking
from .paragraph import ParagraphDocumentChunking
def get_document_chunker(chunking_strategy: str):
if chunking_strategy == ChunkingStrategy.LAYOUT.value:
return LayoutDocumentChunking()
elif chunking_strategy == ChunkingStrategy.PAGE.value:
return PageDocumentChunking()
elif chunking_strategy == ChunkingStrategy.FIXED_SIZE_OVERLAP.value:
return FixedSizeOverlapDocumentChunking()
elif chunking_strategy == ChunkingStrategy.PARAGRAPH.value:
return ParagraphDocumentChunking()
else:
raise Exception(f"Unknown chunking strategy: {chunking_strategy}")