101 lines
2.9 KiB
Python
101 lines
2.9 KiB
Python
 
|
"""Core data models for document chunking."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from enum import Enum
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from pydantic import BaseModel, Field
|
||
|
|
|
||
|
|
|
||
|
|
class BlockType(str, Enum):
|
||
|
|
HEADING = "heading"
|
||
|
|
PARAGRAPH = "paragraph"
|
||
|
|
TABLE = "table"
|
||
|
|
IMAGE = "image"
|
||
|
|
LIST = "list"
|
||
|
|
CODE = "code"
|
||
|
|
|
||
|
|
|
||
|
|
class SplitMode(str, Enum):
|
||
|
|
DEFAULT = "default"
|
||
|
|
DELIMITER = "delimiter"
|
||
|
|
PARENT_CHILD = "parent_child"
|
||
|
|
BY_ROW = "by_row"
|
||
|
|
|
||
|
|
|
||
|
|
class Block(BaseModel):
|
||
|
|
"""Atomic document unit; images and tables must not be split across chunks."""
|
||
|
|
|
||
|
|
type: BlockType
|
||
|
|
text: str = ""
|
||
|
|
markdown: str = ""
|
||
|
|
level: int = 0
|
||
|
|
image_id: str | None = None
|
||
|
|
image_path: str | None = None
|
||
|
|
ocr_text: str = ""
|
||
|
|
meta: dict[str, Any] = Field(default_factory=dict)
|
||
|
|
|
||
|
|
def render(self) -> str:
|
||
|
|
if self.type == BlockType.TABLE:
|
||
|
|
parts: list[str] = []
|
||
|
|
title = self.meta.get("table_title")
|
||
|
|
if title:
|
||
|
|
parts.append(f"### {title}")
|
||
|
|
image_ref = self.image_path or self.meta.get("crop_path")
|
||
|
|
if image_ref:
|
||
|
|
alt = (title or self.image_id or "table").replace("\n", " ").strip()
|
||
|
|
parts.append(f"")
|
||
|
|
parts.append(self.markdown or self.text)
|
||
|
|
footnotes = self.meta.get("footnotes")
|
||
|
|
if footnotes and footnotes not in (self.markdown or ""):
|
||
|
|
parts.append(f"*{footnotes}*")
|
||
|
|
if self.ocr_text:
|
||
|
|
parts.append(f"*[表格 OCR]* {self.ocr_text}")
|
||
|
|
return "\n\n".join(p for p in parts if p.strip())
|
||
|
|
if self.type == BlockType.IMAGE:
|
||
|
|
alt = (self.image_id or self.text or "image").replace("\n", " ").strip()
|
||
|
|
path = self.image_path or self.image_id or ""
|
||
|
|
body = f""
|
||
|
|
if self.ocr_text:
|
||
|
|
body += f"\n\n*[OCR]* {self.ocr_text}"
|
||
|
|
return body
|
||
|
|
if self.type == BlockType.HEADING:
|
||
|
|
prefix = "#" * max(1, min(self.level, 6))
|
||
|
|
return f"{prefix} {self.text}".strip()
|
||
|
|
return self.text or self.markdown
|
||
|
|
|
||
|
|
|
||
|
|
class SplitConfig(BaseModel):
|
||
|
|
mode: SplitMode = SplitMode.DEFAULT
|
||
|
|
delimiter: str | None = None
|
||
|
|
parent_delimiter: str | None = None
|
||
|
|
child_delimiter: str | None = None
|
||
|
|
max_chunk_size: int = 1500
|
||
|
|
child_max_size: int = 512
|
||
|
|
overlap: int = 150
|
||
|
|
header_row_start: int = 1
|
||
|
|
header_row_end: int = 1
|
||
|
|
start_row: int = 2
|
||
|
|
rows_per_chunk: int = 1
|
||
|
|
table_format: str = "markdown"
|
||
|
|
|
||
|
|
|
||
|
|
class Chunk(BaseModel):
|
||
|
|
index: int
|
||
|
|
content: str
|
||
|
|
char_count: int
|
||
|
|
block_types: list[str]
|
||
|
|
meta: dict[str, Any] = Field(default_factory=dict)
|
||
|
|
|
||
|
|
|
||
|
|
class ChunkResult(BaseModel):
|
||
|
|
filename: str
|
||
|
|
doc_id: str
|
||
|
|
split_mode: SplitMode
|
||
|
|
split_config: dict[str, Any] = Field(default_factory=dict)
|
||
|
|
block_count: int
|
||
|
|
chunk_count: int
|
||
|
|
chunks: list[Chunk]
|
||
|
|
assets_dir: str | None = None
|