110 lines
3.3 KiB
Python
110 lines
3.3 KiB
Python
from dataclasses import dataclass
|
|
from typing import Optional
|
|
from config import settings
|
|
from logger import get_logger
|
|
|
|
log = get_logger("chunker")
|
|
|
|
|
|
@dataclass
|
|
class Chunk:
|
|
text: str
|
|
index: int
|
|
source: str
|
|
page: Optional[int] = None
|
|
sheet: Optional[str] = None
|
|
metadata: Optional[dict] = None
|
|
|
|
|
|
class TextChunker:
|
|
def __init__(
|
|
self,
|
|
chunk_size: int = None,
|
|
overlap: int = None
|
|
):
|
|
self.chunk_size = chunk_size or settings.CHUNK_SIZE
|
|
self.overlap = overlap or settings.CHUNK_OVERLAP
|
|
log.debug(f"TextChunker: size={self.chunk_size}, overlap={self.overlap}")
|
|
|
|
def chunk_text(
|
|
self,
|
|
text: str,
|
|
source: str,
|
|
page: Optional[int] = None,
|
|
sheet: Optional[str] = None,
|
|
metadata: Optional[dict] = None
|
|
) -> list[Chunk]:
|
|
words = text.split()
|
|
log.debug(f"Chunking text: {len(words)} words, source={source}, page={page}, sheet={sheet}")
|
|
|
|
if len(words) <= self.chunk_size:
|
|
log.debug(f"Text fits in single chunk")
|
|
return [Chunk(
|
|
text=text,
|
|
index=0,
|
|
source=source,
|
|
page=page,
|
|
sheet=sheet,
|
|
metadata=metadata
|
|
)]
|
|
|
|
chunks = []
|
|
start = 0
|
|
idx = 0
|
|
|
|
while start < len(words):
|
|
end = start + self.chunk_size
|
|
chunk_words = words[start:end]
|
|
chunk_text = " ".join(chunk_words)
|
|
|
|
chunks.append(Chunk(
|
|
text=chunk_text,
|
|
index=idx,
|
|
source=source,
|
|
page=page,
|
|
sheet=sheet,
|
|
metadata=metadata
|
|
))
|
|
|
|
log.debug(f"Chunk {idx}: {len(chunk_words)} words, pos {start}-{end}")
|
|
start = end - self.overlap
|
|
idx += 1
|
|
|
|
log.debug(f"Created {len(chunks)} chunks")
|
|
return chunks
|
|
|
|
def chunk_document(self, doc: dict) -> list[Chunk]:
|
|
log.info(f"Chunking document: {doc.get('source')}")
|
|
all_chunks = []
|
|
|
|
if "pages" in doc:
|
|
log.debug(f"Document has {len(doc['pages'])} pages/sheets")
|
|
for page_data in doc["pages"]:
|
|
page_chunks = self.chunk_text(
|
|
text=page_data["text"],
|
|
source=doc["source"],
|
|
page=page_data.get("page"),
|
|
sheet=page_data.get("sheet"),
|
|
metadata={"type": doc["type"]}
|
|
)
|
|
all_chunks.extend(page_chunks)
|
|
if page_data.get("page"):
|
|
log.debug(f"Page {page_data['page']}: {len(page_chunks)} chunks")
|
|
elif page_data.get("sheet"):
|
|
log.debug(f"Sheet '{page_data['sheet']}': {len(page_chunks)} chunks")
|
|
else:
|
|
log.debug(f"Single page document")
|
|
all_chunks = self.chunk_text(
|
|
text=doc["full_text"],
|
|
source=doc["source"],
|
|
metadata={"type": doc["type"]}
|
|
)
|
|
|
|
log.info(f"Document chunked: {len(all_chunks)} total chunks")
|
|
return all_chunks
|
|
|
|
|
|
def chunk_document(doc: dict) -> list[Chunk]:
|
|
chunker = TextChunker()
|
|
return chunker.chunk_document(doc)
|