Mahdi Bazrafshan
8715d9c186
Why: - Core evaluation targets — each strategy chunks differently for comparison Changes: - fixed_size: token-based splitting with overlap (baseline) - recursive: cascade splitting (headers → newlines → sentences → words) - semantic: sentence-level embeddings with similarity-based boundaries - contextual_retrieval: LLM-generated context prefixes per chunk (Anthropic research) - semantic_parent_child: paragraph clustering into parent-child hierarchy
123 lines
3.8 KiB
Python
123 lines
3.8 KiB
Python
"""Semantic chunking strategy.
|
|
|
|
Sentence-level granularity (ADR 0012):
|
|
1. Split markdown into sentences.
|
|
2. Embed each sentence via OpenAI.
|
|
3. Compute cosine similarity between adjacent sentences.
|
|
4. When similarity drops below SEMANTIC_THRESHOLD, create a chunk boundary.
|
|
5. Enforce SEMANTIC_MIN_CHUNK_SIZE (minimum sentences per chunk).
|
|
6. Boundary sentence stays with the previous chunk.
|
|
|
|
Note: this strategy requires embeddings at chunk-time. The chunk()
|
|
method returns text chunks WITHOUT embeddings — the embedding step
|
|
happens in the orchestration layer (service.py) which calls the
|
|
embedding service after chunking.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from src.chunking.base import ChunkingStrategy, build_chunk, split_sentences
|
|
from src.core.config import settings
|
|
from src.core.models import Chunk, DocumentTree, StrategyName
|
|
|
|
|
|
def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
|
|
"""Compute cosine similarity between two vectors."""
|
|
dot = np.dot(a, b)
|
|
norm = np.linalg.norm(a) * np.linalg.norm(b)
|
|
if norm == 0:
|
|
return 0.0
|
|
return float(dot / norm)
|
|
|
|
|
|
def _group_sentences_into_chunks(
|
|
sentences: list[str],
|
|
embeddings: list[list[float]],
|
|
threshold: float,
|
|
min_size: int,
|
|
) -> list[str]:
|
|
"""Group sentences into chunks based on semantic similarity.
|
|
|
|
Returns a list of chunk texts.
|
|
"""
|
|
if not sentences:
|
|
return []
|
|
if len(sentences) <= min_size:
|
|
return [" ".join(sentences)]
|
|
|
|
chunks: list[str] = []
|
|
current_group: list[str] = [sentences[0]]
|
|
|
|
for i in range(1, len(sentences)):
|
|
sim = _cosine_similarity(
|
|
np.array(embeddings[i - 1]),
|
|
np.array(embeddings[i]),
|
|
)
|
|
|
|
if sim < threshold and len(current_group) >= min_size:
|
|
# Topic shift — close current chunk
|
|
chunks.append(" ".join(current_group))
|
|
current_group = [sentences[i]]
|
|
else:
|
|
current_group.append(sentences[i])
|
|
|
|
# Flush remaining
|
|
if current_group:
|
|
# If the last group is too small, merge into previous
|
|
if chunks and len(current_group) < min_size:
|
|
last = chunks.pop()
|
|
chunks.append(last + " " + " ".join(current_group))
|
|
else:
|
|
chunks.append(" ".join(current_group))
|
|
|
|
return chunks
|
|
|
|
|
|
class SemanticStrategy(ChunkingStrategy):
|
|
name = StrategyName.SEMANTIC
|
|
|
|
def chunk(
|
|
self,
|
|
*,
|
|
doc_name: str,
|
|
tree: DocumentTree,
|
|
markdown: str,
|
|
sentence_embeddings: list[list[float]] | None = None,
|
|
) -> list[Chunk]:
|
|
"""Produce semantic chunks.
|
|
|
|
If sentence_embeddings is provided (from the orchestration layer),
|
|
uses them for boundary detection. Otherwise, falls back to
|
|
paragraph-level chunking (sentences without similarity-based splits).
|
|
"""
|
|
sentences = split_sentences(markdown)
|
|
if not sentences:
|
|
return []
|
|
|
|
threshold = settings.semantic_threshold
|
|
min_size = settings.semantic_min_chunk_size
|
|
|
|
if sentence_embeddings and len(sentence_embeddings) == len(sentences):
|
|
chunk_texts = _group_sentences_into_chunks(
|
|
sentences, sentence_embeddings, threshold, min_size
|
|
)
|
|
else:
|
|
# Fallback: group sentences into fixed-size chunks
|
|
chunk_texts = []
|
|
for i in range(0, len(sentences), min_size):
|
|
group = sentences[i:i + min_size]
|
|
chunk_texts.append(" ".join(group))
|
|
|
|
chunks: list[Chunk] = []
|
|
for i, text in enumerate(chunk_texts):
|
|
if text.strip():
|
|
chunks.append(build_chunk(
|
|
strategy=self.name,
|
|
doc_name=doc_name,
|
|
index=i,
|
|
text=text.strip(),
|
|
))
|
|
return chunks
|