Spaces:
Sleeping
Sleeping
| import hashlib | |
| import re | |
| from collections.abc import Iterable | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from techdocs_graphrag.models import Chunk | |
| SUPPORTED_SUFFIXES = frozenset({".md", ".markdown", ".txt"}) | |
| PARAGRAPH_SEPARATOR = "\n\n" | |
| class SourceDocument: | |
| name: str | |
| text: str | |
| def load_documents(root: Path) -> list[SourceDocument]: | |
| return [ | |
| SourceDocument(name=path.relative_to(root).as_posix(), text=path.read_text(encoding="utf-8")) | |
| for path in sorted(root.rglob("*")) | |
| if path.is_file() and path.suffix.lower() in SUPPORTED_SUFFIXES | |
| ] | |
| class ParagraphSplitter: | |
| def __init__(self, chunk_size: int, chunk_overlap: int) -> None: | |
| if not 0 <= chunk_overlap < chunk_size: | |
| raise ValueError("chunk_overlap must be in [0, chunk_size)") | |
| self._chunk_size = chunk_size | |
| self._chunk_overlap = chunk_overlap | |
| def split(self, text: str) -> list[str]: | |
| chunks: list[str] = [] | |
| buffer = "" | |
| for paragraph in self._paragraphs(text): | |
| candidate = self._join(buffer, paragraph) | |
| if len(candidate) <= self._chunk_size: | |
| buffer = candidate | |
| continue | |
| if buffer: | |
| chunks.append(buffer) | |
| candidate = self._join(self._overlap_of(buffer), paragraph) | |
| while len(candidate) > self._chunk_size: | |
| chunks.append(candidate[: self._chunk_size]) | |
| candidate = candidate[self._chunk_size - self._chunk_overlap :] | |
| buffer = candidate | |
| if buffer: | |
| chunks.append(buffer) | |
| return chunks | |
| def _paragraphs(text: str) -> Iterable[str]: | |
| return (block.strip() for block in re.split(r"\n\s*\n", text) if block.strip()) | |
| def _join(head: str, tail: str) -> str: | |
| return f"{head}{PARAGRAPH_SEPARATOR}{tail}" if head else tail | |
| def _overlap_of(self, text: str) -> str: | |
| if not self._chunk_overlap or len(text) <= self._chunk_overlap: | |
| return "" | |
| tail = text[-self._chunk_overlap :] | |
| return tail.partition(" ")[2] or tail | |
| def chunk_id(document: str, position: int, text: str) -> str: | |
| return hashlib.sha256(f"{document}\x00{position}\x00{text}".encode()).hexdigest()[:16] | |
| def chunk_documents(documents: Iterable[SourceDocument], splitter: ParagraphSplitter) -> list[Chunk]: | |
| return [ | |
| Chunk(id=chunk_id(document.name, position, text), document=document.name, position=position, text=text) | |
| for document in documents | |
| for position, text in enumerate(splitter.split(document.text)) | |
| ] | |