Spaces:
Running
Running
| from __future__ import annotations | |
| import asyncio | |
| import logging | |
| import re | |
| from typing import List | |
| logger = logging.getLogger(__name__) | |
| def chunk_text( | |
| text: str, | |
| chunk_size: int = 512, | |
| chunk_overlap: int = 64, | |
| ) -> List[str]: | |
| if chunk_overlap >= chunk_size: | |
| chunk_overlap = chunk_size // 4 | |
| paragraphs = re.split(r"\n\s*\n", text.strip()) | |
| chunks: List[str] = [] | |
| current: List[str] = [] | |
| current_len = 0 | |
| for para in paragraphs: | |
| para = para.strip() | |
| if not para: | |
| continue | |
| para_len = len(para) | |
| if current_len + para_len + 1 <= chunk_size: | |
| current.append(para) | |
| current_len += para_len + 1 | |
| else: | |
| if current: | |
| chunks.append("\n\n".join(current)) | |
| if para_len > chunk_size: | |
| for i in range(0, para_len, chunk_size): | |
| segment = para[i:i + chunk_size] | |
| if len(segment) >= chunk_size // 3: | |
| chunks.append(segment) | |
| current = [] | |
| current_len = 0 | |
| else: | |
| current = [para] | |
| current_len = para_len + 1 | |
| if current: | |
| chunks.append("\n\n".join(current)) | |
| if chunk_overlap > 0 and len(chunks) > 1: | |
| overlapped: List[str] = [] | |
| for i, chunk in enumerate(chunks): | |
| if i == 0: | |
| overlapped.append(chunk) | |
| else: | |
| prev = chunks[i - 1] | |
| overlap_text = " ".join(prev.split()[-chunk_overlap:]) if len(prev.split()) > chunk_overlap else prev | |
| combined = f"{overlap_text}\n\n{chunk}" | |
| overlapped.append(combined) | |
| return overlapped if all(len(c) <= chunk_size + chunk_overlap + 10 for c in overlapped) else chunks | |
| return chunks if chunks else [text] | |
| async def chunk_text_async( | |
| text: str, | |
| chunk_size: int = 512, | |
| chunk_overlap: int = 64, | |
| ) -> List[str]: | |
| loop = asyncio.get_running_loop() | |
| return await loop.run_in_executor(None, chunk_text, text, chunk_size, chunk_overlap) | |