62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
|
|
from ocr_pipeline.config import settings
|
|
from ocr_pipeline.utils.logging import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class TextChunk:
|
|
content: str
|
|
chunk_index: int
|
|
start_char: int
|
|
end_char: int
|
|
metadata: dict[str, Any]
|
|
|
|
|
|
class SemanticChunker:
|
|
def __init__(self):
|
|
self.splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=settings.chunking.chunk_size,
|
|
chunk_overlap=settings.chunking.chunk_overlap,
|
|
separators=settings.chunking.separators,
|
|
keep_separator=settings.chunking.keep_separator,
|
|
length_function=len,
|
|
)
|
|
|
|
def chunk(self, text: str, base_metadata: dict[str, Any] | None = None) -> list[TextChunk]:
|
|
if not text or not text.strip():
|
|
return []
|
|
|
|
base_metadata = base_metadata or {}
|
|
docs = self.splitter.create_documents([text], metadatas=[base_metadata])
|
|
|
|
chunks = []
|
|
for i, doc in enumerate(docs):
|
|
start = text.find(doc.page_content)
|
|
if start == -1:
|
|
start = 0
|
|
chunks.append(
|
|
TextChunk(
|
|
content=doc.page_content,
|
|
chunk_index=i,
|
|
start_char=start,
|
|
end_char=start + len(doc.page_content),
|
|
metadata={**doc.metadata, "chunk_index": i, "total_chunks": len(docs)},
|
|
)
|
|
)
|
|
|
|
logger.debug("chunked_text", chunks=len(chunks), original_length=len(text))
|
|
return chunks
|
|
|
|
|
|
def chunk_text(text: str, metadata: dict[str, Any] | None = None) -> list[TextChunk]:
|
|
chunker = SemanticChunker()
|
|
return chunker.chunk(text, metadata)
|