feat(strategies): implement all 5 chunking strategies
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
This commit is contained in:
122
src/chunking/strategies/semantic.py
Normal file
122
src/chunking/strategies/semantic.py
Normal file
@@ -0,0 +1,122 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user