""" Document chunking strategies for ResearchPilot. Three strategies implemented: 1. FixedSizeChunker — baseline, educational 2. RecursiveChunker — production standard 3. SemanticChunker — highest quality, used in final system Each chunker produces a list of Chunk objects with identical structure so the rest of the pipeline doesn't care which strategy was used. This is the STRATEGY PATTERN in software design. """ import re import uuid import json from pathlib import Path from dataclasses import dataclass, field, asdict from typing import Optional import logging logging.getLogger("sentence_transformers").setLevel(logging.ERROR) import numpy as np from langchain_text_splitters import RecursiveCharacterTextSplitter from src.utils.logger import get_logger from config.settings import ( CHUNK_SIZE, CHUNK_OVERLAP, MIN_CHUNK_SIZE, CHUNKS_DIR, EMBEDDING_MODEL_NAME ) logger = get_logger(__name__) # --------------------------------------------------------- # DATA MODEL # --------------------------------------------------------- @dataclass class Chunk: """ A single chunk of text with all metadata attached. WHY ATTACH METADATA TO EVERY CHUNK: When a user asks a question and we retrieve chunk #347, we need to know: which paper did this come from? What page? What section? Without metadata on the chunk itself, we'd have to do a separate lookup — slow and error-prone. Every chunk is self-contained and self-describing. """ # Unique identifier for this chunk # uuid4() generates a random unique ID - no two chunks collide chunk_id: str # The actual text content text: str # Which paper this came from paper_id: str title: str authors: list[str] published_date: str primary_category: str arxiv_url: str # Position within the document chunk_index: int # 0, 1, 2, ...(position in paper) total_chunks: int # How many chunks this paper was split into # Text statistics char_count: int word_count: int # Chunking metadata chunking_strategy: str # 'fixed', 'recursive', 'semantic' def to_dict(self) -> dict: """Convert to dict for JSON serialization""" return asdict(self) @property def is_valid(self) -> bool: """Check if chunk has enough content to be useful.""" return ( len(self.text.strip()) >= MIN_CHUNK_SIZE and self.word_count >= 10 # At least 10 words ) # --------------------------------------------------------- # STRATEGY 1: FIXED SIZE CHUNKER # --------------------------------------------------------- class FixedSizeChunker: """ Splits text every N characters with M character overlap. This is the WORST chunking strategy but we include it as: 1. A baseline to compare against 2. To demonstrate WHY better strategies exist 3. Educational — see exactly what breaks OVERLAP EXPLAINED: Without overlap: Chunk 1: "The model achieves 94.2% accuracy on" Chunk 2: "GLUE benchmark, beating prior work by" The phrase "accuracy on GLUE" is split — neither chunk contains the complete concept. With 50-char overlap: Chunk 1: "The model achieves 94.2% accuracy on" Chunk 2: "accuracy on GLUE benchmark, beating prior work by" Now "accuracy on GLUE" exists in chunk 2. Retrieval works. Overlap is a band-aid for fixed-size chunking's core problem. """ def __init__( self, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP ): self.chunk_size = chunk_size self.overlap = overlap def split(self, text: str, metadata: dict) -> list[Chunk]: chunks = [] start = 0 index = 0 while start < len(text): end = min(start + self.chunk_size, len(text)) chunk_text = text[start : end].strip() if len(chunk_text) > MIN_CHUNK_SIZE: chunks.append(self._make_chunk(chunk_text, index, metadata)) index += 1 # Move forward by (chunk_size - overlap) # This creates the sliding window effect start += self.chunk_size - self.overlap # Now that we know total_chunks, update all chunks for chunk in chunks: chunk.total_chunks = len(chunks) return chunks def _make_chunk(self, text: str, index: int, meta: dict) -> Chunk: return Chunk( chunk_id = str(uuid.uuid4()), text = text, paper_id = meta['paper_id'], title = meta['title'], authors = meta['authors'], published_date = meta['published_date'], primary_category = meta['primary_category'], arxiv_url = meta['arxiv_url'], chunk_index = index, total_chunks = 0, # Updated after all chunks created char_count = len(text), word_count = len(text.split()), chunking_strategy = 'fixed', ) # --------------------------------------------------------- # STRATEGY 2: RECURSIVE CHARACTER SPLITTER # --------------------------------------------------------- class RecursiveChunker: """ Splits text by trying delimiters in order of preference. DELIMITER HIERARCHY: 1. "\n\n" -> paragraph break (best — complete thought) 2. "\n" -> line break (good) 3. ". " -> sentence end (acceptable) 4. " " -> word boundary (last resort) 5. "" -> character (never want this) The splitter tries to split at \n\n first. If a resulting piece is still too large, it tries \n. Still too large? Tries ". " etc. WHY THIS IS BETTER THAN FIXED: Fixed chunking: "...achieves 94.2% ac" + "curacy on GLUE..." Recursive: "...achieves 94.2% accuracy on GLUE benchmark." Recursive splitting respects natural language boundaries. The resulting chunks contain complete sentences and paragraphs. THIS IS THE INDUSTRY STANDARD. Use this unless you have a specific reason to use semantic chunking. """ def __init__( self, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP ): # LangChain's implementation is well-tested and handles # many edge cases we'd miss writing our own self.splitter = RecursiveCharacterTextSplitter( chunk_size = chunk_size, chunk_overlap = overlap, length_function = len, # Separators tried in order - most preferred first separators = ["\n\n", "\n", ". ", " ", ""], # Keep separator at the end of the chunk (preserves sentence endings) keep_separator = True, ) def split(self, text: str, metadata: dict) -> list[Chunk]: # LangChain splits the text into string pieces text_pieces = self.splitter.split_text(text) chunks = [] for index, piece in enumerate(text_pieces): piece = piece.strip() if len(piece) < MIN_CHUNK_SIZE: continue chunk = Chunk( chunk_id = str(uuid.uuid4()), text = piece, paper_id = metadata["paper_id"], title = metadata["title"], authors = metadata["authors"], published_date = metadata["published_date"], primary_category = metadata["primary_category"], arxiv_url = metadata["arxiv_url"], chunk_index = index, total_chunks = 0, char_count = len(piece), word_count = len(piece.split()), chunking_strategy = "recursive", ) chunks.append(chunk) for chunk in chunks: chunk.total_chunks = len(chunks) return chunks # --------------------------------------------------------- # STRATEGY 3: SEMANTIC CHUNKER # --------------------------------------------------------- class SemanticChunker: """ Splits text at points where the semantic meaning changes. THE CORE INSIGHT: In a research paper, adjacent sentences that discuss the SAME idea have HIGH embedding similarity. When the topic shifts (e.g., from "method" to "results"), the similarity between adjacent sentences DROPS sharply. We find these DROP POINTS and split there. ALGORITHM: 1. Split text into individual sentences 2. Embed every sentence using BGE model 3. Calculate cosine similarity between each adjacent pair: sim(sentence_1, sentence_2), sim(sentence_2, sentence_3), ... 4. Find similarity values that drop below a threshold (these are semantic boundaries) 5. Split the document at those boundary points 6. Each resulting chunk contains sentences about ONE topic VISUAL EXAMPLE: Sentence similarities: [0.92, 0.89, 0.91, 0.45, 0.88, 0.90, 0.38, 0.85] ↑ split here ↑ split here The drops at 0.45 and 0.38 mark topic changes. WHY THIS MATTERS FOR RESEARCH PAPERS: Papers have clear sections: Introduction -> Method -> Experiments -> Conclusion Within each section, sentences are semantically close. At section transitions, similarity drops sharply. Semantic chunking naturally aligns with paper structure. """ def __init__( self, model_name: str = EMBEDDING_MODEL_NAME, # Similarity Threshold: splits happen where similarity < this value # 0.5 means "splits when adjacent sentences share < 50% semantic similarity" breakpoint_threshold: float = 0.5, # Minimum sentences per chunk (avoid 1-sentence chunks) min_sentences_per_chunk: int = 3, # Maximum sentences per chunk (avoid enormous chunks) max_sentences_per_chunk: int = 15, ): self.breakpoint_threshold = breakpoint_threshold self.min_sentences_per_chunk = min_sentences_per_chunk self.max_sentences_per_chunk = max_sentences_per_chunk # Lazy load the model - only load when first needed # WHY: Loading a 110MB model takes ~3 seconds # We don't want that delay at import time self._model = None self._model_name = model_name logger.info( f"SemanticChunker initialized " f"(threshold={breakpoint_threshold})" ) @property def model(self): """ Lazy-load the embedding model. PROPERTY PATTERN: self.model looks like an attribute but actually runs this function the first time it's accessed. After first load, self._model is set and returned directly. """ if self._model is None: from sentence_transformers import SentenceTransformer logger.info(f"Loading embedding model: {self._model_name}") self._model = SentenceTransformer(self._model_name) logger.info("Embedding model loaded") return self._model def _split_into_sentences(self, text: str) -> list[str]: """ Split text into individual sentences. WHY NOT JUST split('.'): "Dr. Smith proposed..." -> split on period gives ["Dr", " Smith proposed"] "The accuracy was 94.2% on..." -> breaks at decimal "et al. showed..." -> breaks at abbreviation Our regex handles these cases: - Requires capital letter after period (new sentence starts with capital) - Handles "." followed by newline - Keeps sentences of reasonable length """ # Split on: period/!/? followed by whitespace and capital letter # OR on double newlines (paragraph breaks are always sentence breaks) sentence_pattern = r'(?<=[.|?])\s+(?=[A-Z])|(?<=\n)\n+' sentences = re.split(sentence_pattern, text) # Filter out very short fragments (less than 20 chars) # These are usually artifacts, not real sentences sentences = [s.strip() for s in sentences if len(s.strip()) > 20] return sentences def _compute_similarities(self, sentences: list[str]) -> np.ndarray: """ Embed all sentences and compute cosine similarity between each adjacent pair. Returns array of shape (len(sentences) - 1,) where result[i] = cosine_similarity(sentence[i], sentence[i+1]) BATCH PROCESSING: We embed ALL sentences at once (not one by one). Batch embedding is 10-50x faster than individual embedding because the model processes multiple inputs in parallel. """ if len(sentences) < 2: return np.array([]) # encode() returns numpy array of shape (n_sentences, embedding_dim) # show_progress_bar = False keeps output clean in pipeline embeddings = self.model.encode( sentences, show_progress_bar = False, batch_size = 64, # Process 64 sentences at a time normalize_embeddings = True, # L2 Normalize for cosine similarity convert_to_numpy = True, # Skip torch tensor, go direct to numpy ) # Vectorized similarity — faster than Python loop # embeddings[:-1] = all sentences except last # embeddings[1:] = all sentences except first # Row-wise dot product = cosine similarity for normalized vectors similarities = np.sum( embeddings[:-1] * embeddings[1:], axis = 1 ) return similarities def _find_split_points(self, similarities: np.ndarray) -> list[int]: """ Find indices where the text should be split. ADAPTIVE THRESHOLD: Instead of a fixed threshold, we use the mean - std_dev. This adapts to each document's similarity distribution. WHY ADAPTIVE: A technical paper might have overall lower similarities than a narrative paper. A fixed threshold of 0.5 might split a technical paper every sentence (too granular) while never splitting a narrative paper (too coarse). Adaptive threshold = "split where similarity is significantly below average for THIS document." """ if len(similarities) == 0: return [] # Adaptive threshold: mean minus one standard deviation # This finds the "unusually low similarity" points mean_sim = np.mean(similarities) std_min = np.std(similarities) threshold = max( self.breakpoint_threshold, # Never go above configured max mean_sim - std_min # Adaptive: 1 std below mean ) logger.debug( f"Similarity stats: mean={mean_sim:.3f}" f"std={std_min:.3f}, threshold={threshold:.3f}" ) # Find indices where similarity drops below threshold # These are the semantic breakpoints split_points = [ i + 1 # +1 because we split AFTER sentence 1 for i, sim in enumerate(similarities) if sim < threshold ] return split_points def _group_sentences_into_chunks( self, sentences: list[str], split_points: list[int] ) -> list[str]: """ Group sentences into chunks based on split points, respecting min/max sentence constraints. """ if not sentences: return [] # Build groups using split_points as boundaries # split_points = [4, 9, 15] means: # Group 1: sentences 0-3 # Group 2: sentences 4-8 # Group 3: sentences 9-14 # Group 4: sentences 15+ boundaries = [0] + split_points + [len(sentences)] groups = [] for i in range(len(boundaries) - 1): start = boundaries[i] end = boundaries[i + 1] group = sentences[start : end] if not group: continue # ENFORCE MINIMUM: If group is too small, merge with next if len(group) < self.min_sentences_per_chunk and group and groups: # Merge into previous group groups[-1].extend(group) else: groups.append(group) # ENFORCE MAXIMUM: If group is too large, subdivide it final_group = [] for group in groups: if len(group) <= self.max_sentences_per_chunk: final_group.append(group) else: # Split large groups into max_size pieces for i in range(0, len(group), self.max_sentences_per_chunk): sub = group[i : i + self.max_sentences_per_chunk] if sub: final_group.append(sub) # Convert sentence lists to text strings return [" ".join(group) for group in final_group] def split(self, text: str, metadata: dict) -> list[Chunk]: """ Main split method - full semantic chunking pipeline """ # Step 1: Split into sentences sentences = self._split_into_sentences(text) if len(sentences) < 2: # Documents too short for semantic analysis - fall back to recursive logger.debug( f"Too few sentences ({len(sentences)}) for semantic" f"chunking on {metadata['paper_id']}, using recursive" ) return RecursiveChunker().split(text, metadata) # Step 2: Compute inter-sentence similarities similarities = self._compute_similarities(sentences) # Step 3: Find semantic breakpoints split_points = self._find_split_points(similarities) logger.debug( f"{metadata['paper_id']}: {len(sentences)} sentences, " f"{len(split_points)} splits points found" ) # Step 4: Group sentences into chunks chunk_texts = self._group_sentences_into_chunks(sentences, split_points) # Step 5: Build Chunk objects chunks = [] for index, chunk_text in enumerate(chunk_texts): chunk_text = chunk_text.strip() if len(chunk_text) < MIN_CHUNK_SIZE: continue chunk = Chunk( chunk_id = str(uuid.uuid4()), text = chunk_text, paper_id = metadata["paper_id"], title = metadata["title"], authors = metadata["authors"], published_date = metadata["published_date"], primary_category = metadata["primary_category"], arxiv_url = metadata["arxiv_url"], chunk_index = index, total_chunks = 0, char_count = len(chunk_text), word_count = len(chunk_text.split()), chunking_strategy = "semantic", ) chunks.append(chunk) for chunk in chunks: chunk.total_chunks = len(chunks) logger.debug( f"{metadata['paper_id']}: produced {len(chunks)} semantic chunks" ) return chunks # --------------------------------------------------------- # PIPELINE RUNNER # --------------------------------------------------------- class ChunkingPipeline: """ Orchestrates chunking for all processed papers. Takes files from data/processed/ and produces chunk files in data/chunks/. """ def __init__(self, strategy: str = 'recursive'): """ Args: strategy: "fixed" | "recursive" | "semantic" """ valid = {"fixed", "recursive", "semantic"} if strategy not in valid: raise ValueError(f"Strategy must be one of {valid}") self.strategy_name = strategy # Instantiate the correct chunker if strategy == "fixed": self.chunker = FixedSizeChunker() elif strategy == "recursive": self.chunker = RecursiveChunker() elif strategy == "semantic": self.chunker = SemanticChunker() logger.info(f"ChunkingPipeline initialized with strategy: {strategy}") def process_paper(self, processed_doc: dict) -> list[Chunk]: """Chunk a single processed document""" paper_id = processed_doc['paper_id'] text = processed_doc.get("full_text", "") if not text: logger.warning(f"No text found for {paper_id}") return [] # Metadata dict passes to every chunk metadata = { "paper_id": paper_id, "title": processed_doc.get("title", ""), "authors": processed_doc.get("authors", []), "published_date": processed_doc.get("published_date", ""), "primary_category": processed_doc.get("primary_category") or (processed_doc.get("categories") or ["cs.LG"])[0], "arxiv_url": processed_doc.get("arxiv_url", ""), } return self.chunker.split(text, metadata) def save_chunks(self, chunks: list[Chunk], paper_id: str): """ Save all chunks for a paper to data/chunks/. File format: data/chunks/{paper_id}_{strategy}.json Contains list of chunk dicts. """ if not chunks: return output_path = ( CHUNKS_DIR / f"{paper_id}_{self.strategy_name}.json" ) with open(output_path, 'w', encoding = 'utf-8') as f: json.dump( [chunk.to_dict() for chunk in chunks], f, indent = 2, ensure_ascii = False ) def run(self, process_dir: Path) -> dict: """ Run chunking pipeline on all processed documents. Args: processed_dir: Path to data/processed/ Returns: Summary statistics """ from tqdm import tqdm processed_files = list(process_dir.glob("*.json")) logger.info( f"Chunking {len(processed_files)} documents " f"with '{self.strategy_name}' strategy" ) total_chunks = 0 successful = 0 failed = 0 skipped = 0 for proc_file in tqdm(processed_files, desc = f"Chunking ({self.strategy_name})"): with open(proc_file, 'r', encoding = 'utf-8') as f: doc = json.load(f) paper_id = doc['paper_id'] output_path = CHUNKS_DIR / f"{paper_id}_{self.strategy_name}.json" # Skip already chunked (idempotent) if output_path.exists(): skipped += 1 continue try: chunks = self.process_paper(doc) if not chunks: failed += 1 continue self.save_chunks(chunks, paper_id) total_chunks += len(chunks) successful += 1 except Exception as e: logger.error(f"Failed to chunk {paper_id}: {e}") failed += 1 stats = { "strategy": self.strategy_name, "documents": len(processed_files), "successful": successful, "failed": failed, "skipped": skipped, "total_chunks": total_chunks, "avg_chunks_per_doc": ( round(total_chunks / max(successful, 1), 1) ), } logger.info(f"Chunking complete: {stats}") return stats