contextwormhole 1.0.0

Teleport beyond context limits with transformers

ContextWormhole

Context length extension library for transformers

ContextWormhole provides practical implementations of three established context extension techniques. When your transformer model reaches its context limit, this library offers clean, tested strategies to handle longer inputs.

pip install contextwormhole

Purpose

Most transformer models have fixed context windows (e.g., 1024 tokens for GPT-2). This library implements three strategies to work with longer texts while maintaining the model's original architecture.

Strategies

1. Sliding Window

Processes text in overlapping chunks, maintaining continuity between segments.

@sliding_window(window_size=512, overlap=64)
def process_long_document(model, text, **kwargs):
    return model.generate(text, **kwargs)

• Implementation: Overlapping windows with position ID recycling
• Time complexity: O(n)
• Memory complexity: O(window_size)
• Use cases: Documents, code files, articles

2. Hierarchical Context

Creates summaries of text chunks, then combines summaries with final content.

@hierarchical_context(chunk_size=256, summary_length=64)
def analyze_paper(model, paper, **kwargs):
    return model.generate(paper, **kwargs)

• Implementation: Chunk → summarize → combine → process
• Time complexity: O(n log n)
• Memory complexity: O(n/chunk_size * summary_length)
• Use cases: Research papers, structured documents

3. Attention Sink

Preserves initial tokens plus recent context, discarding middle content.

@attention_sink(sink_tokens=16)
def continue_conversation(model, chat_history, **kwargs):
    return model.generate(chat_history, **kwargs)

• Implementation: Initial tokens + recent context
• Time complexity: O(1)
• Memory complexity: O(max_length)
• Use cases: Conversations, chat histories

Empirical Results

Tests on repetition patterns (10 runs each, distilgpt2):

Strategy	Uniqueness Ratio	Repeated Phrases	Notes
Standard (low temp)	0.59	3.5	Baseline
Standard (high temp)	0.28	2.0	High repetition
Attention Sink	0.67	1.8	Best coherence

The attention sink strategy showed consistently better text quality with fewer repetitive patterns.

Usage

Basic Example

from contextwormhole import ContextWormholeModel

model = ContextWormholeModel("gpt2")

# Different strategies for different needs
result1 = model.sliding_window_generate(long_document, max_new_tokens=100)
result2 = model.hierarchical_generate(research_paper, max_new_tokens=100)
result3 = model.attention_sink_generate(conversation_history, max_new_tokens=100)

Configuration

from contextwormhole import ExtendedContextConfig

config = ExtendedContextConfig(
    window_size=256,
    overlap=64,
    chunk_size=256,
    summary_length=64,
    sink_tokens=16,
    use_cache=True,
)

model = ContextWormholeModel("gpt2", **config.__dict__)

CLI Interface

# Sliding window
contextwormhole --model gpt2 --input document.txt --strategy sliding_window

# Hierarchical
contextwormhole --model gpt2 --input paper.txt --strategy hierarchical

# Attention sink
contextwormhole --model gpt2 --input chat.txt --strategy attention_sink

Performance Characteristics

Strategy	Max Context	Memory (MB)*	Time (s)*	Best For
Sliding Window	~10K tokens	600	1.5-2.0	Documents, code
Hierarchical	~20K tokens	400	1.0-1.5	Papers, reports
Attention Sink	~8K tokens	300	0.8-1.2	Conversations

*Approximate values for GPT-2 on CPU

Implementation Notes

• Each strategy respects the model's native context limit for individual forward passes
• Position ID recycling enables handling of arbitrarily long inputs
• KV caching improves generation speed and maintains coherence
• All strategies include proper error handling and configuration validation

Requirements

• Python ≥ 3.8
• PyTorch ≥ 1.9.0
• Transformers ≥ 4.20.0
• NumPy ≥ 1.20.0

Technical Background

This library implements well-established context extension techniques:

• Sliding Window: Classical attention windowing
• Hierarchical Context: Recursive summarization approach
• Attention Sink: Based on StreamingLLM research

The focus is on providing clean, tested implementations with practical optimizations rather than novel algorithms.

License

MIT License