Spaces:
Running
Running
File size: 2,702 Bytes
5301c48 |
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 |
from typing import List
import re
from starfish.data_ingest.splitter.base_splitter import TextSplitter
class TextChunkSplitter(TextSplitter):
"""Splitting text into chunks with optional overlap and minimum size constraints."""
def __init__(
self,
chunk_size: int = 400,
overlap: int = 20,
min_chunk_size: int = 100,
) -> None:
"""Create a new TextChunkSplitter.
Args:
chunk_size: Maximum size of each chunk
overlap: Number of characters to overlap between chunks
min_chunk_size: Minimum acceptable chunk size (avoids tiny final chunks)
"""
self._chunk_size = chunk_size
self._overlap = overlap
self._min_chunk_size = min_chunk_size
def split_text(self, text: str) -> List[str]:
"""Split text into chunks with optional overlap.
Args:
text: Input text to split
Returns:
List of text chunks
"""
# Normalize whitespace and handle different paragraph separators
text = re.sub(r"\n{2,}", "\n\n", text.strip())
paragraphs = text.split("\n\n")
chunks = []
current_chunk = ""
for para in paragraphs:
# Skip empty paragraphs
if not para.strip():
continue
# If adding this paragraph would exceed chunk size
if current_chunk and len(current_chunk) + len(para) > self._chunk_size:
# Ensure we don't create chunks smaller than min_chunk_size
if len(current_chunk) >= self._min_chunk_size:
chunks.append(current_chunk)
# Create overlap using sentence boundaries
sentences = [s for s in re.split(r"(?<=[.!?])\s+", current_chunk) if s]
overlap_text = ""
# Add sentences until we reach the desired overlap
for sentence in reversed(sentences):
if len(overlap_text) + len(sentence) <= self._overlap:
overlap_text = sentence + " " + overlap_text
else:
break
current_chunk = overlap_text.strip() + "\n\n" + para
else:
# If chunk is too small, keep adding to it
current_chunk += "\n\n" + para
else:
current_chunk += ("\n\n" + para) if current_chunk else para
# Add the final chunk if it meets minimum size
if current_chunk and len(current_chunk) >= self._min_chunk_size:
chunks.append(current_chunk)
return chunks
|