Spaces:
Sleeping
Sleeping
File size: 2,700 Bytes
da21448 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | 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"
@dataclass(frozen=True, slots=True)
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
@staticmethod
def _paragraphs(text: str) -> Iterable[str]:
return (block.strip() for block in re.split(r"\n\s*\n", text) if block.strip())
@staticmethod
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))
]
|