File size: 5,117 Bytes
f3997d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
"""
Text chunking utilities for document processing.
"""
from typing import List
import re


class TextChunker:
    """Handles intelligent text chunking with various strategies."""
    
    @staticmethod
    def chunk_by_sentences(
        text: str,
        chunk_size: int = 800,
        overlap: int = 200
    ) -> List[str]:
        """
        Chunk text by sentences with overlap.
        
        Args:
            text: Text to chunk
            chunk_size: Target size of each chunk in characters
            overlap: Overlap between chunks in characters
            
        Returns:
            List of text chunks
        """
        if not text or len(text.strip()) == 0:
            return []
        
        # Split into sentences (improved regex for better sentence detection)
        sentences = re.split(r'(?<=[.!?])\s+', text)
        
        chunks = []
        current_chunk = []
        current_size = 0
        
        for sentence in sentences:
            sentence_size = len(sentence)
            
            # If adding this sentence exceeds chunk_size, save current chunk
            if current_size + sentence_size > chunk_size and current_chunk:
                chunk_text = ' '.join(current_chunk)
                chunks.append(chunk_text)
                
                # Calculate overlap: keep last few sentences
                overlap_text = []
                overlap_size = 0
                for s in reversed(current_chunk):
                    if overlap_size + len(s) <= overlap:
                        overlap_text.insert(0, s)
                        overlap_size += len(s)
                    else:
                        break
                
                current_chunk = overlap_text
                current_size = overlap_size
            
            current_chunk.append(sentence)
            current_size += sentence_size
        
        # Add remaining chunk
        if current_chunk:
            chunks.append(' '.join(current_chunk))
        
        return [c.strip() for c in chunks if c.strip()]
    
    @staticmethod
    def chunk_by_paragraphs(
        text: str,
        max_chunk_size: int = 1000
    ) -> List[str]:
        """
        Chunk text by paragraphs, combining small paragraphs.
        
        Args:
            text: Text to chunk
            max_chunk_size: Maximum size of each chunk
            
        Returns:
            List of text chunks
        """
        if not text or len(text.strip()) == 0:
            return []
        
        # Split by double newlines (paragraphs)
        paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]
        
        chunks = []
        current_chunk = []
        current_size = 0
        
        for para in paragraphs:
            para_size = len(para)
            
            # If paragraph alone exceeds max size, split it by sentences
            if para_size > max_chunk_size:
                # Save current chunk if exists
                if current_chunk:
                    chunks.append('\n\n'.join(current_chunk))
                    current_chunk = []
                    current_size = 0
                
                # Split large paragraph by sentences
                sentence_chunks = TextChunker.chunk_by_sentences(
                    para,
                    chunk_size=max_chunk_size,
                    overlap=100
                )
                chunks.extend(sentence_chunks)
                continue
            
            # If adding this paragraph exceeds max size, save current chunk
            if current_size + para_size > max_chunk_size and current_chunk:
                chunks.append('\n\n'.join(current_chunk))
                current_chunk = []
                current_size = 0
            
            current_chunk.append(para)
            current_size += para_size + 2  # +2 for \n\n
        
        # Add remaining chunk
        if current_chunk:
            chunks.append('\n\n'.join(current_chunk))
        
        return [c.strip() for c in chunks if c.strip()]
    
    @staticmethod
    def chunk_with_metadata(
        text: str,
        chunk_size: int = 800,
        overlap: int = 200,
        strategy: str = "sentences"
    ) -> List[dict]:
        """
        Chunk text and return with metadata.
        
        Args:
            text: Text to chunk
            chunk_size: Target chunk size
            overlap: Overlap size
            strategy: Chunking strategy ("sentences" or "paragraphs")
            
        Returns:
            List of dictionaries with chunk text and metadata
        """
        if strategy == "paragraphs":
            chunks = TextChunker.chunk_by_paragraphs(text, chunk_size)
        else:
            chunks = TextChunker.chunk_by_sentences(text, chunk_size, overlap)
        
        return [
            {
                "text": chunk,
                "index": i,
                "size": len(chunk),
                "strategy": strategy
            }
            for i, chunk in enumerate(chunks)
        ]


# Global chunker instance
text_chunker = TextChunker()