Spaces:
Runtime error
Runtime error
File size: 7,539 Bytes
45c5a08 | 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 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | import re
from typing import List, Dict, Any
import tiktoken
class ContentChunker:
def __init__(self, max_tokens: int = 512):
"""
Initialize the content chunker with max tokens per chunk
"""
self.max_tokens = max_tokens
self.tokenizer = tiktoken.get_encoding("cl100k_base") # Good for most text
def count_tokens(self, text: str) -> int:
"""
Count the number of tokens in a text
"""
return len(self.tokenizer.encode(text))
def chunk_text(self, text: str, source_path: str = "", chunk_id_prefix: str = "") -> List[Dict[str, Any]]:
"""
Chunk text into segments of max_tokens or less
"""
if not text.strip():
return []
# Split text into sentences to avoid cutting in the middle of sentences
sentences = re.split(r'(?<=[.!?])\s+', text)
chunks = []
current_chunk = ""
current_token_count = 0
chunk_index = 0
for sentence in sentences:
sentence_token_count = self.count_tokens(sentence)
# If a single sentence is too long, we need to break it down
if sentence_token_count > self.max_tokens:
# Split the long sentence into smaller parts
sub_chunks = self._split_long_sentence(sentence)
for sub_chunk in sub_chunks:
sub_chunk_token_count = self.count_tokens(sub_chunk)
if current_token_count + sub_chunk_token_count > self.max_tokens and current_chunk:
# Save current chunk and start new one
chunk_id = f"{chunk_id_prefix}_{chunk_index}" if chunk_id_prefix else str(chunk_index)
chunks.append({
"id": chunk_id,
"content": current_chunk.strip(),
"token_count": current_token_count,
"source_path": source_path,
"chunk_index": chunk_index
})
current_chunk = sub_chunk
current_token_count = sub_chunk_token_count
chunk_index += 1
else:
# Add to current chunk
if current_chunk:
current_chunk += " " + sub_chunk
else:
current_chunk = sub_chunk
current_token_count += sub_chunk_token_count
else:
# Check if adding this sentence would exceed the limit
if current_token_count + sentence_token_count > self.max_tokens and current_chunk:
# Save current chunk and start new one
chunk_id = f"{chunk_id_prefix}_{chunk_index}" if chunk_id_prefix else str(chunk_index)
chunks.append({
"id": chunk_id,
"content": current_chunk.strip(),
"token_count": current_token_count,
"source_path": source_path,
"chunk_index": chunk_index
})
current_chunk = sentence
current_token_count = sentence_token_count
chunk_index += 1
else:
# Add sentence to current chunk
if current_chunk:
current_chunk += " " + sentence
else:
current_chunk = sentence
current_token_count += sentence_token_count
# Add the last chunk if it has content
if current_chunk.strip():
chunk_id = f"{chunk_id_prefix}_{chunk_index}" if chunk_id_prefix else str(chunk_index)
chunks.append({
"id": chunk_id,
"content": current_chunk.strip(),
"token_count": current_token_count,
"source_path": source_path,
"chunk_index": chunk_index
})
return chunks
def _split_long_sentence(self, sentence: str) -> List[str]:
"""
Split a sentence that is too long into smaller parts
"""
if self.count_tokens(sentence) <= self.max_tokens:
return [sentence]
# Try to split by commas first
parts = sentence.split(', ')
if all(self.count_tokens(part) <= self.max_tokens for part in parts):
return [part.strip() + ', ' if i < len(parts) - 1 else part.strip()
for i, part in enumerate(parts)]
# If comma splitting doesn't work, split by words
words = sentence.split()
chunks = []
current_chunk = ""
for word in words:
test_chunk = current_chunk + " " + word if current_chunk else word
if self.count_tokens(test_chunk) <= self.max_tokens:
current_chunk = test_chunk
else:
if current_chunk: # If there's something to save
chunks.append(current_chunk.strip())
current_chunk = word
if current_chunk: # Add the last chunk
chunks.append(current_chunk.strip())
return chunks
def chunk_markdown(self, markdown_content: str, source_path: str = "") -> List[Dict[str, Any]]:
"""
Chunk markdown content preserving section structure where possible
"""
# Split by markdown headers to keep sections together when possible
header_pattern = r'^(#{1,6})\s+(.+)$'
lines = markdown_content.split('\n')
sections = []
current_section = {'header': '', 'content': '', 'level': 0}
for line in lines:
header_match = re.match(header_pattern, line.strip())
if header_match:
# Save current section if it has content
if current_section['content'].strip():
sections.append({
'header': current_section['header'],
'content': current_section['content'].strip(),
'level': current_section['level']
})
# Start new section
header_level = len(header_match.group(1))
header_text = header_match.group(2)
current_section = {
'header': header_text,
'content': f"{'#' * header_level} {header_text}\n\n",
'level': header_level
}
else:
current_section['content'] += line + '\n'
# Add the last section
if current_section['content'].strip():
sections.append({
'header': current_section['header'],
'content': current_section['content'].strip(),
'level': current_section['level']
})
# Now chunk each section
all_chunks = []
for i, section in enumerate(sections):
section_content = section['content']
section_chunks = self.chunk_text(section_content, source_path, f"section_{i}")
# Add section metadata to each chunk
for chunk in section_chunks:
chunk['section_header'] = section['header']
chunk['section_level'] = section['level']
all_chunks.append(chunk)
return all_chunks |