Initial scaffold: .gitignore, config, requirements, data seeds, directory structure
This commit is contained in:
n0mad1k
@@ -0,0 +1,253 @@
|
||||
"""Text chunking, cleaning, and token estimation for Mosaic."""
|
||||
import re
|
||||
import hashlib
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Rough token estimation: ~4 chars per token for English text
|
||||
CHARS_PER_TOKEN = 4
|
||||
|
||||
# Default chunk sizes
|
||||
DEFAULT_CHUNK_TOKENS = 8000
|
||||
DEFAULT_OVERLAP_TOKENS = 500
|
||||
EMAIL_MAX_TOKENS = 4000
|
||||
PDF_MAX_TOKENS = 8000
|
||||
|
||||
|
||||
def estimate_tokens(text: str) -> int:
|
||||
"""Estimate token count from text length."""
|
||||
return len(text) // CHARS_PER_TOKEN
|
||||
|
||||
|
||||
def clean_text(text: str) -> str:
|
||||
"""Clean extracted text: normalize whitespace, remove control chars."""
|
||||
if not text:
|
||||
return ""
|
||||
# Remove null bytes and other control characters (keep newlines, tabs)
|
||||
text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)
|
||||
# Normalize line endings
|
||||
text = text.replace('\r\n', '\n').replace('\r', '\n')
|
||||
# Collapse excessive blank lines (3+ → 2)
|
||||
text = re.sub(r'\n{3,}', '\n\n', text)
|
||||
# Strip trailing whitespace per line
|
||||
text = '\n'.join(line.rstrip() for line in text.split('\n'))
|
||||
return text.strip()
|
||||
|
||||
|
||||
def chunk_text(text: str, doc_type: str = "text",
|
||||
max_tokens: int = DEFAULT_CHUNK_TOKENS,
|
||||
overlap_tokens: int = DEFAULT_OVERLAP_TOKENS,
|
||||
metadata: Optional[dict] = None) -> list[dict]:
|
||||
"""Chunk text using per-document-type strategy.
|
||||
|
||||
Args:
|
||||
text: The text to chunk
|
||||
doc_type: Document type (html, pdf, email, cable, text)
|
||||
max_tokens: Maximum tokens per chunk
|
||||
overlap_tokens: Token overlap between chunks
|
||||
metadata: Optional metadata (e.g., cable headers)
|
||||
|
||||
Returns:
|
||||
List of dicts with 'text', 'chunk_index', 'total_chunks', 'metadata'
|
||||
"""
|
||||
if not text:
|
||||
return []
|
||||
|
||||
if doc_type == 'cable':
|
||||
return _chunk_cable(text, max_tokens, overlap_tokens, metadata)
|
||||
elif doc_type == 'pdf':
|
||||
return _chunk_pdf(text, max_tokens, overlap_tokens)
|
||||
elif doc_type == 'email':
|
||||
return _chunk_email(text, max_tokens, overlap_tokens, metadata)
|
||||
elif doc_type == 'html':
|
||||
return _chunk_html(text, max_tokens, overlap_tokens)
|
||||
else:
|
||||
return _chunk_default(text, max_tokens, overlap_tokens)
|
||||
|
||||
|
||||
def _chunk_cable(text: str, max_tokens: int, overlap_tokens: int,
|
||||
metadata: Optional[dict] = None) -> list[dict]:
|
||||
"""Chunk cable: headers as structured context, body chunked independently."""
|
||||
headers = metadata or {}
|
||||
header_text = ""
|
||||
if headers:
|
||||
header_lines = []
|
||||
for key in ('SUBJECT', 'ORIGIN', 'CLASSIFICATION', 'TAGS', 'DATE'):
|
||||
if key in headers:
|
||||
header_lines.append(f"{key}: {headers[key]}")
|
||||
header_text = '\n'.join(header_lines)
|
||||
|
||||
# If small enough, return as one chunk
|
||||
total_tokens = estimate_tokens(header_text + '\n\n' + text)
|
||||
if total_tokens <= max_tokens:
|
||||
return [{'text': (header_text + '\n\n' + text).strip(),
|
||||
'chunk_index': 0, 'total_chunks': 1, 'metadata': headers}]
|
||||
|
||||
# Chunk body, prepend headers to each chunk
|
||||
body_chunks = _chunk_default(text, max_tokens - estimate_tokens(header_text) - 50,
|
||||
overlap_tokens)
|
||||
for i, chunk in enumerate(body_chunks):
|
||||
chunk['text'] = header_text + '\n\n' + chunk['text']
|
||||
chunk['metadata'] = headers
|
||||
chunk['chunk_index'] = i
|
||||
chunk['total_chunks'] = len(body_chunks)
|
||||
return body_chunks
|
||||
|
||||
|
||||
def _chunk_pdf(text: str, max_tokens: int, overlap_tokens: int) -> list[dict]:
|
||||
"""Chunk PDF by page boundaries (indicated by form feeds or page markers)."""
|
||||
# Split by form feed or common page break markers
|
||||
pages = re.split(r'\f|\n-{3,}\n|\n={3,}\n', text)
|
||||
pages = [p.strip() for p in pages if p.strip()]
|
||||
|
||||
if not pages:
|
||||
return _chunk_default(text, max_tokens, overlap_tokens)
|
||||
|
||||
chunks = []
|
||||
current_text = ""
|
||||
for page in pages:
|
||||
combined_tokens = estimate_tokens(current_text + '\n\n' + page)
|
||||
if combined_tokens > max_tokens and current_text:
|
||||
chunks.append(current_text.strip())
|
||||
# Keep overlap from end of previous chunk
|
||||
overlap_chars = overlap_tokens * CHARS_PER_TOKEN
|
||||
current_text = current_text[-overlap_chars:] + '\n\n' + page
|
||||
else:
|
||||
current_text = (current_text + '\n\n' + page).strip()
|
||||
|
||||
if current_text.strip():
|
||||
chunks.append(current_text.strip())
|
||||
|
||||
return [{'text': c, 'chunk_index': i, 'total_chunks': len(chunks), 'metadata': {}}
|
||||
for i, c in enumerate(chunks)]
|
||||
|
||||
|
||||
def _chunk_email(text: str, max_tokens: int, overlap_tokens: int,
|
||||
metadata: Optional[dict] = None) -> list[dict]:
|
||||
"""Chunk email: headers + body as atomic unit unless body is too large."""
|
||||
headers = metadata or {}
|
||||
total_tokens = estimate_tokens(text)
|
||||
|
||||
if total_tokens <= EMAIL_MAX_TOKENS:
|
||||
return [{'text': text, 'chunk_index': 0, 'total_chunks': 1, 'metadata': headers}]
|
||||
|
||||
# Split headers from body
|
||||
header_section = ""
|
||||
body = text
|
||||
if '\n\n' in text:
|
||||
parts = text.split('\n\n', 1)
|
||||
header_section = parts[0]
|
||||
body = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
# Chunk body, prepend headers
|
||||
body_chunks = _chunk_default(body, max_tokens - estimate_tokens(header_section) - 50,
|
||||
overlap_tokens)
|
||||
for i, chunk in enumerate(body_chunks):
|
||||
chunk['text'] = header_section + '\n\n' + chunk['text']
|
||||
chunk['metadata'] = headers
|
||||
chunk['chunk_index'] = i
|
||||
chunk['total_chunks'] = len(body_chunks)
|
||||
return body_chunks
|
||||
|
||||
|
||||
def _chunk_html(text: str, max_tokens: int, overlap_tokens: int) -> list[dict]:
|
||||
"""Chunk HTML: section-aware chunking preserving heading hierarchy."""
|
||||
# Split on headings
|
||||
sections = re.split(r'(?=\n#{1,6}\s|\n[A-Z][A-Za-z\s]{3,}\n[=-]+\n)', text)
|
||||
sections = [s for s in sections if s.strip()]
|
||||
|
||||
if not sections or len(sections) == 1:
|
||||
return _chunk_default(text, max_tokens, overlap_tokens)
|
||||
|
||||
chunks = []
|
||||
current_text = ""
|
||||
for section in sections:
|
||||
combined_tokens = estimate_tokens(current_text + '\n\n' + section)
|
||||
if combined_tokens > max_tokens and current_text:
|
||||
chunks.append(current_text.strip())
|
||||
overlap_chars = overlap_tokens * CHARS_PER_TOKEN
|
||||
current_text = current_text[-overlap_chars:] + '\n\n' + section
|
||||
else:
|
||||
current_text = (current_text + '\n\n' + section).strip()
|
||||
|
||||
if current_text.strip():
|
||||
chunks.append(current_text.strip())
|
||||
|
||||
return [{'text': c, 'chunk_index': i, 'total_chunks': len(chunks), 'metadata': {}}
|
||||
for i, c in enumerate(chunks)]
|
||||
|
||||
|
||||
def _chunk_default(text: str, max_tokens: int, overlap_tokens: int) -> list[dict]:
|
||||
"""Default chunking: paragraph-aware with token overlap."""
|
||||
total_tokens = estimate_tokens(text)
|
||||
if total_tokens <= max_tokens:
|
||||
return [{'text': text, 'chunk_index': 0, 'total_chunks': 1, 'metadata': {}}]
|
||||
|
||||
max_chars = max_tokens * CHARS_PER_TOKEN
|
||||
overlap_chars = overlap_tokens * CHARS_PER_TOKEN
|
||||
|
||||
# Split into paragraphs
|
||||
paragraphs = text.split('\n\n')
|
||||
chunks = []
|
||||
current = ""
|
||||
|
||||
for para in paragraphs:
|
||||
if len(current) + len(para) + 2 > max_chars and current:
|
||||
chunks.append(current.strip())
|
||||
# Keep overlap
|
||||
current = current[-overlap_chars:] + '\n\n' + para
|
||||
else:
|
||||
current = (current + '\n\n' + para) if current else para
|
||||
|
||||
if current.strip():
|
||||
chunks.append(current.strip())
|
||||
|
||||
# If still too large (single giant paragraph), force-split
|
||||
final_chunks = []
|
||||
for chunk in chunks:
|
||||
if len(chunk) > max_chars * 1.5:
|
||||
# Force split at sentence boundaries or character limit
|
||||
for sub in _force_split(chunk, max_chars, overlap_chars):
|
||||
final_chunks.append(sub)
|
||||
else:
|
||||
final_chunks.append(chunk)
|
||||
|
||||
return [{'text': c, 'chunk_index': i, 'total_chunks': len(final_chunks), 'metadata': {}}
|
||||
for i, c in enumerate(final_chunks)]
|
||||
|
||||
|
||||
def _force_split(text: str, max_chars: int, overlap_chars: int) -> list[str]:
|
||||
"""Force-split text at sentence boundaries or character limits."""
|
||||
sentences = re.split(r'(?<=[.!?])\s+', text)
|
||||
chunks = []
|
||||
current = ""
|
||||
|
||||
for sent in sentences:
|
||||
if len(current) + len(sent) + 1 > max_chars and current:
|
||||
chunks.append(current.strip())
|
||||
current = current[-overlap_chars:] + ' ' + sent
|
||||
else:
|
||||
current = (current + ' ' + sent) if current else sent
|
||||
|
||||
if current.strip():
|
||||
chunks.append(current.strip())
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
def dedup_text(texts: list[str], threshold: float = 0.95) -> list[str]:
|
||||
"""Remove near-duplicate texts using hash comparison.
|
||||
|
||||
For exact duplicates. Near-duplicate detection would require
|
||||
more sophisticated approaches (minhash, simhash, etc.)
|
||||
"""
|
||||
seen = set()
|
||||
unique = []
|
||||
for text in texts:
|
||||
h = hashlib.md5(text.encode()).hexdigest()
|
||||
if h not in seen:
|
||||
seen.add(h)
|
||||
unique.append(text)
|
||||
return unique
|
||||
Reference in New Issue
Block a user