Spaces:
Paused
Paused
| """Document parsing and configurable chunking. | |
| Documents are markdown-ish text with `## section` headings. The parser splits | |
| into sections (keeping section titles as metadata), and the chunker splits | |
| long sections into overlapping windows measured in words. | |
| """ | |
| from __future__ import annotations | |
| import re | |
| from typing import Dict, List | |
| from pydantic import BaseModel | |
| class Chunk(BaseModel): | |
| chunk_id: str | |
| text: str | |
| metadata: Dict[str, str] | |
| def parse_sections(text: str) -> List[Dict[str, str]]: | |
| """Split a document into {section, text} blocks on '## ' headings.""" | |
| sections: List[Dict[str, str]] = [] | |
| current_title = "preamble" | |
| current_lines: List[str] = [] | |
| for line in text.splitlines(): | |
| m = re.match(r"^##\s+(.*)", line) | |
| if m: | |
| if current_lines and "".join(current_lines).strip(): | |
| sections.append({"section": current_title, | |
| "text": "\n".join(current_lines).strip()}) | |
| current_title = m.group(1).strip() | |
| current_lines = [] | |
| else: | |
| current_lines.append(line) | |
| if current_lines and "".join(current_lines).strip(): | |
| sections.append({"section": current_title, | |
| "text": "\n".join(current_lines).strip()}) | |
| return sections | |
| def chunk_words(text: str, chunk_size: int = 160, overlap: int = 40) -> List[str]: | |
| """Overlapping word-window chunking. chunk_size/overlap are in words.""" | |
| if chunk_size <= 0: | |
| raise ValueError("chunk_size must be positive") | |
| if overlap >= chunk_size: | |
| raise ValueError("overlap must be smaller than chunk_size") | |
| words = text.split() | |
| if len(words) <= chunk_size: | |
| return [text.strip()] if text.strip() else [] | |
| step = chunk_size - overlap | |
| chunks = [] | |
| for start in range(0, len(words), step): | |
| window = words[start:start + chunk_size] | |
| if not window: | |
| break | |
| chunks.append(" ".join(window)) | |
| if start + chunk_size >= len(words): | |
| break | |
| return chunks | |
| def chunk_document(doc_id: str, text: str, base_metadata: Dict[str, str], | |
| chunk_size: int = 160, overlap: int = 40) -> List[Chunk]: | |
| """Parse into sections then chunk each section, propagating metadata.""" | |
| chunks: List[Chunk] = [] | |
| for si, sec in enumerate(parse_sections(text)): | |
| for ci, piece in enumerate(chunk_words(sec["text"], chunk_size, overlap)): | |
| meta = dict(base_metadata) | |
| meta["section"] = sec["section"] | |
| chunks.append(Chunk(chunk_id=f"{doc_id}:s{si}:c{ci}", text=piece, | |
| metadata=meta)) | |
| return chunks | |