Context Window Management: Handling Long Documents in LLM Applications

Working with large language models (LLMs) requires careful management of context windows due to token limitations. This technical guide explores practical strategies for handling long documents while maintaining context coherence and processing efficiency. We'll cover implementation approaches with concrete examples and specific optimization techniques.

Understanding Token Limits

Token Basics

LLMs process text as tokens, with most models having fixed context window sizes:

GPT-3.5: 4,096 tokens
GPT-4: 8,192 or 32,768 tokens
Claude: 100,000 tokens

One token typically represents 4 characters in English, though this varies by language and content type. Special tokens like code snippets or numbers may encode differently.

Impact on Processing

Example token count estimation
def estimate_tokens(text: str) -> int:
    return len(text.split()) * 1.3  # Rough approximation

Chunking Strategies

Basic Document Chunking

Divide documents into manageable segments while preserving semantic coherence:

def chunk_document(text: str, chunk_size: int = 2000) -> List[str]:
    sentences = nltk.sent_tokenize(text)
    chunks = []
    current_chunk = []
    current_size = 0
    
    for sentence in sentences:
        sentence_size = len(sentence.split())
        if current_size + sentence_size > chunk_size:
            chunks.append(' '.join(current_chunk))
            current_chunk = [sentence]
            current_size = sentence_size
        else:
            current_chunk.append(sentence)
            current_size += sentence_size
            
    if current_chunk:
        chunks.append(' '.join(current_chunk))
    return chunks

Semantic Chunking

More sophisticated approach using semantic boundaries:

from transformers import AutoTokenizer, AutoModel
def semantic_chunk(text: str, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name)
    
    # Implementation details for semantic segmentation
    # Based on sentence embeddings and similarity scores

Sliding Window Techniques

Overlap Implementation

Maintain context continuity between chunks:

def sliding_window(text: str, window_size: int = 2000, overlap: int = 200):
    chunks = []
    start = 0
    text_length = len(text)
    
    while start < text_length:
        end = min(start + window_size, text_length)
        
        # Adjust end to complete last word
        if end < text_length:
            end = text.rfind(' ', start, end)
            
        chunk = text[start:end]
        chunks.append(chunk)
        
        start = end - overlap
    return chunks

Context Preservation

Maintain relevant information across chunks:

def preserve_context(chunks: List[str], context_size: int = 100):
    enhanced_chunks = []
    for i, chunk in enumerate(chunks):
        context = ""
        if i > 0:
            context = chunks[i-1][-context_size:]
        enhanced_chunks.append(f"{context}\n{chunk}")
    return enhanced_chunks

When to Use RAG

RAG vs. Chunking Decision Matrix

Consider these factors when choosing between chunking and RAG:

1. Document Size - < 10K tokens: Direct processing - 10K-100K tokens: Chunking - > 100K tokens: RAG

2. Query Pattern - Single-pass analysis: Chunking - Multiple queries: RAG - Real-time updates: RAG

RAG Implementation Example

from langchain import Document, VectorStore
from langchain.embeddings import OpenAIEmbeddings
def setup_rag(documents: List[str]):
    embeddings = OpenAIEmbeddings()
    texts = [Document(page_content=doc) for doc in documents]
    vector_store = VectorStore.from_documents(texts, embeddings)
    return vector_store

Performance Optimization

Parallel Processing

Implement concurrent chunk processing:

from concurrent.futures import ThreadPoolExecutor
def process_chunks_parallel(chunks: List[str], max_workers: int = 4):
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [executor.submit(process_chunk, chunk) for chunk in chunks]
        results = [f.result() for f in futures]
    return results

Memory Management

Optimize memory usage during processing:

def optimize_memory(chunks: List[str]):
    # Generator-based processing
    for chunk in chunks:
        processed = process_chunk(chunk)
        yield processed
        del chunk  # Explicit cleanup

Real-World Use Cases

Document Analysis Pipeline

class DocumentProcessor:
    def __init__(self, chunk_size: int = 2000, overlap: int = 200):
        self.chunk_size = chunk_size
        self.overlap = overlap
        
    def process_document(self, document_path: str):
        text = self.load_document(document_path)
        chunks = sliding_window(text, self.chunk_size, self.overlap)
        results = process_chunks_parallel(chunks)
        return self.aggregate_results(results)

Large-Scale Text Processing

Example implementation for processing multiple documents:

def batch_process_documents(doc_paths: List[str], batch_size: int = 10):
    processor = DocumentProcessor()
    results = []
    
    for i in range(0, len(doc_paths), batch_size):
        batch = doc_paths[i:i+batch_size]
        batch_results = [processor.process_document(path) for path in batch]
        results.extend(batch_results)
    
    return results

Key Takeaways

1. Choose chunking strategy based on document characteristics and processing requirements 2. Implement overlap to maintain context continuity 3. Consider RAG for very large documents or multiple query scenarios 4. Optimize performance through parallel processing and memory management 5. Monitor and adjust chunk sizes based on model limitations and performance metrics

Next Steps

Implement error handling and recovery mechanisms
Add logging and monitoring
Develop automated testing for chunk processing
Consider implementing a caching layer for frequently accessed chunks
Evaluate and benchmark different chunking strategies for your specific use case