contextwormhole 1.1.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

Why Position ID Recycling?

Position IDs are critical in transformer models as they provide information about token order. However, they present a significant challenge when working with inputs that exceed the model's maximum context length:

1. Index Out of Range Errors: Without proper handling, position IDs for long inputs can exceed the maximum index in the position embedding table, causing runtime errors.

2. Context Preservation: Simply truncating inputs loses valuable context. Position ID recycling allows us to maintain more context by intelligently selecting which parts of the input to keep.

3. Quality Improvements: Our tests show that proper position ID handling reduces repetition in generated text and improves overall coherence.

4. Arbitrary Length Handling: With position ID recycling, the library can process inputs of any length while ensuring position IDs always stay within the valid range (0 to max_position_embeddings-1).

The implementation uses modulo arithmetic to "recycle" position IDs, combined with strategic token selection to preserve the most relevant context from beginning, middle, and end of long documents.

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