Spaces:
Running
Running
File size: 3,128 Bytes
5ea3240 | 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 79 80 81 82 83 84 85 86 87 88 89 90 | 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,
)
|