Spaces:
Running
Running
| from __future__ import annotations | |
| import re | |
| import uuid | |
| import numpy as np | |
| from .config import get_settings | |
| from .schemas import Chunk, Document | |
| from .security import prompt_injection_score | |
| def _split_sentences(text: str) -> list[str]: | |
| text = re.sub(r"\s+", " ", text).strip() | |
| if not text: | |
| return [] | |
| return re.split(r"(?<=[.!?])\s+(?=[A-Z0-9])", text) | |
| def chunk_documents(documents: list[Document], semantic: bool = False) -> list[Chunk]: | |
| settings = get_settings() | |
| chunks: list[Chunk] = [] | |
| for doc in documents: | |
| sentences = _split_sentences(doc.text) | |
| if semantic and len(sentences) >= 4: | |
| sentences = _semantic_groups(sentences) | |
| if not sentences: | |
| continue | |
| current: list[str] = [] | |
| current_len = 0 | |
| for sentence in sentences: | |
| if current and current_len + len(sentence) + 1 > settings.chunk_size_chars: | |
| text = " ".join(current).strip() | |
| chunks.append(_make_chunk(doc, text)) | |
| overlap: list[str] = [] | |
| overlap_len = 0 | |
| for item in reversed(current): | |
| if overlap_len + len(item) > settings.chunk_overlap_chars: | |
| break | |
| overlap.insert(0, item) | |
| overlap_len += len(item) + 1 | |
| current = overlap | |
| current_len = sum(len(x) + 1 for x in current) | |
| current.append(sentence) | |
| current_len += len(sentence) + 1 | |
| if current: | |
| chunks.append(_make_chunk(doc, " ".join(current).strip())) | |
| return chunks[: settings.max_chunks_per_session] | |
| def _semantic_groups(sentences: list[str]) -> list[str]: | |
| """Group adjacent sentences at semantic breakpoints before size-based chunking.""" | |
| try: | |
| from .retrieval import ModelRegistry | |
| vectors = np.asarray(list(ModelRegistry.embedding().passage_embed(sentences))) | |
| norms = np.linalg.norm(vectors, axis=1, keepdims=True) + 1e-9 | |
| vectors = vectors / norms | |
| sims = np.sum(vectors[:-1] * vectors[1:], axis=1) | |
| threshold = float(np.percentile(sims, 20)) | |
| groups: list[str] = [] | |
| current = [sentences[0]] | |
| current_len = len(sentences[0]) | |
| for i, sentence in enumerate(sentences[1:]): | |
| should_break = sims[i] <= threshold and current_len >= 500 | |
| if should_break: | |
| groups.append(" ".join(current)) | |
| current = [sentence] | |
| current_len = len(sentence) | |
| else: | |
| current.append(sentence) | |
| current_len += len(sentence) + 1 | |
| if current: | |
| groups.append(" ".join(current)) | |
| return groups | |
| except Exception: | |
| return sentences | |
| def _make_chunk(doc: Document, text: str) -> Chunk: | |
| metadata = dict(doc.metadata) | |
| metadata["injection_score"] = prompt_injection_score(text) | |
| return Chunk( | |
| id=str(uuid.uuid4()), | |
| text=text, | |
| source=doc.source, | |
| page=doc.page, | |
| section=doc.section, | |
| metadata=metadata, | |
| ) | |