infinite-context-gateway 0.1.0__tar.gz

infinite_context_gateway-0.1.0/PKG-INFO

@@ -0,0 +1,108 @@
+Metadata-Version: 2.4
+Name: infinite-context-gateway
+Version: 0.1.0
+Summary: A hybrid AI context gateway combining Headroom cloud compression and In-Place Local Test-Time Training (TTT).
+Author-email: Dev Team <dev@example.com>
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: httpx
+Provides-Extra: cloud
+Requires-Dist: headroom-ai; extra == "cloud"
+Provides-Extra: local
+Requires-Dist: torch; extra == "local"
+Requires-Dist: transformers; extra == "local"
+Requires-Dist: peft; extra == "local"
+Requires-Dist: accelerate; extra == "local"
+Requires-Dist: bitsandbytes; extra == "local"
+Provides-Extra: all
+Requires-Dist: headroom-ai; extra == "all"
+Requires-Dist: torch; extra == "all"
+Requires-Dist: transformers; extra == "all"
+Requires-Dist: peft; extra == "all"
+Requires-Dist: accelerate; extra == "all"
+Requires-Dist: bitsandbytes; extra == "all"
+
+# Infinite Context
+
+A hybrid AI context gateway combining **Headroom cloud compression** and **In-Place Local Test-Time Training (TTT)**.
+
+## Installation
+
+To install the base gateway:
+```bash
+pip install infinite_context
+```
+
+To install with Cloud Phase 1 dependencies (Headroom API):
+```bash
+pip install infinite_context[cloud]
+```
+
+To install with Local Phase 2 dependencies (PyTorch, Transformers, PEFT):
+```bash
+pip install infinite_context[local]
+```
+
+## Usage
+
+### Phase 1: Cloud Compression
+Uses the Headroom API to semantically compress a massive context and send it to cloud providers (OpenAI, Anthropic) while keeping costs low.
+
+```python
+from infinite_context import ContextGateway
+
+gateway = ContextGateway(
+    engine="cloud",
+    model_id="claude-3-5-sonnet-20240620",
+    api_key="your_anthropic_api_key",
+    compression_ratio=0.8
+)
+
+# You can pass in conversational history (rolling window)
+history = [
+    {"role": "user", "content": "What is the codebase about?"},
+    {"role": "assistant", "content": "It is a Python application..."}
+]
+
+response = gateway.chat("How does the failover protocol work?", massive_context, history=history)
+print(response)
+```
+
+### Phase 2: Local Test-Time Training (TTT)
+Bypasses the KV-Cache entirely by injecting a PEFT LoRA adapter and baking the context directly into the model's neural weights on your local GPU. Includes Early Stopping latency optimizations and Generation repetition penalties.
+
+```python
+from infinite_context import ContextGateway
+
+gateway = ContextGateway(
+    engine="local",
+    model_id="Qwen/Qwen2.5-0.5B-Instruct",
+    load_in_4bit=True
+)
+
+response = gateway.chat("What is the failover protocol?", massive_context)
+print(response)
+```
+
+### Checkpoint Persistence
+Save trained Fast Weights to disk and resume later without re-reading the document:
+
+```python
+from infinite_context import ContextGateway
+
+# Train and keep state
+gateway = ContextGateway(engine="local", model_id="Qwen/Qwen2.5-0.5B-Instruct")
+gateway.chat("Summarise the document.", massive_context, keep_state=True)
+gateway.save_state("./my_checkpoint")
+
+del gateway  # Free all GPU memory
+
+# Resume later — zero re-training
+gateway = ContextGateway(engine="local", model_id="Qwen/Qwen2.5-0.5B-Instruct")
+gateway.load_state("./my_checkpoint")
+response = gateway.chat("What was the protocol?", context="")
+print(response)
+```

infinite_context_gateway-0.1.0/README.md

@@ -0,0 +1,81 @@
+# Infinite Context
+
+A hybrid AI context gateway combining **Headroom cloud compression** and **In-Place Local Test-Time Training (TTT)**.
+
+## Installation
+
+To install the base gateway:
+```bash
+pip install infinite_context
+```
+
+To install with Cloud Phase 1 dependencies (Headroom API):
+```bash
+pip install infinite_context[cloud]
+```
+
+To install with Local Phase 2 dependencies (PyTorch, Transformers, PEFT):
+```bash
+pip install infinite_context[local]
+```
+
+## Usage
+
+### Phase 1: Cloud Compression
+Uses the Headroom API to semantically compress a massive context and send it to cloud providers (OpenAI, Anthropic) while keeping costs low.
+
+```python
+from infinite_context import ContextGateway
+
+gateway = ContextGateway(
+    engine="cloud",
+    model_id="claude-3-5-sonnet-20240620",
+    api_key="your_anthropic_api_key",
+    compression_ratio=0.8
+)
+
+# You can pass in conversational history (rolling window)
+history = [
+    {"role": "user", "content": "What is the codebase about?"},
+    {"role": "assistant", "content": "It is a Python application..."}
+]
+
+response = gateway.chat("How does the failover protocol work?", massive_context, history=history)
+print(response)
+```
+
+### Phase 2: Local Test-Time Training (TTT)
+Bypasses the KV-Cache entirely by injecting a PEFT LoRA adapter and baking the context directly into the model's neural weights on your local GPU. Includes Early Stopping latency optimizations and Generation repetition penalties.
+
+```python
+from infinite_context import ContextGateway
+
+gateway = ContextGateway(
+    engine="local",
+    model_id="Qwen/Qwen2.5-0.5B-Instruct",
+    load_in_4bit=True
+)
+
+response = gateway.chat("What is the failover protocol?", massive_context)
+print(response)
+```
+
+### Checkpoint Persistence
+Save trained Fast Weights to disk and resume later without re-reading the document:
+
+```python
+from infinite_context import ContextGateway
+
+# Train and keep state
+gateway = ContextGateway(engine="local", model_id="Qwen/Qwen2.5-0.5B-Instruct")
+gateway.chat("Summarise the document.", massive_context, keep_state=True)
+gateway.save_state("./my_checkpoint")
+
+del gateway  # Free all GPU memory
+
+# Resume later — zero re-training
+gateway = ContextGateway(engine="local", model_id="Qwen/Qwen2.5-0.5B-Instruct")
+gateway.load_state("./my_checkpoint")
+response = gateway.chat("What was the protocol?", context="")
+print(response)
+```

infinite_context_gateway-0.1.0/infinite_context/__init__.py

@@ -0,0 +1,10 @@
+"""
+Infinite Context — A hybrid AI context gateway combining Headroom cloud
+compression and In-Place Local Test-Time Training (TTT).
+"""
+
+__version__ = "0.1.0"
+
+from .gateway import ContextGateway
+
+__all__ = ["ContextGateway"]

infinite_context_gateway-0.1.0/infinite_context/cloud/__init__.py

@@ -0,0 +1,6 @@
+"""Cloud engine components — API client and Headroom compression."""
+
+from .client import CloudClient
+from .compressor import HeadroomCompressor
+
+__all__ = ["CloudClient", "HeadroomCompressor"]

infinite_context_gateway-0.1.0/infinite_context/cloud/client.py

@@ -0,0 +1,127 @@
+"""
+Cloud API client for OpenAI, Anthropic, and OpenAI-compatible endpoints.
+
+Handles payload formatting, header construction, and SSE streaming for
+providers like Groq, Together, etc.
+"""
+
+import json
+import httpx
+from typing import Dict, Generator, List
+
+__all__ = ["CloudClient"]
+
+# Shared timeout: 30 s connect, 120 s read (large prompts can be slow).
+_DEFAULT_TIMEOUT = httpx.Timeout(connect=30.0, read=120.0, write=30.0, pool=10.0)
+
+
+class CloudClient:
+    """HTTP client for cloud LLM providers."""
+
+    def __init__(self, api_key: str, model_id: str, base_url: str = None):
+        """
+        Args:
+            api_key: Bearer / x-api-key token.
+            model_id: Model identifier (e.g. ``claude-3-5-sonnet-20240620``).
+            base_url: Custom endpoint URL.  When set, the OpenAI chat-
+                      completions format is assumed.
+        """
+        self.api_key = api_key
+        self.model_id = model_id
+
+        # Determine endpoint and provider format.
+        if base_url:
+            self.base_url = base_url
+            self.provider = "openai"
+        elif "claude" in self.model_id.lower():
+            self.base_url = "https://api.anthropic.com/v1/messages"
+            self.provider = "anthropic"
+        else:
+            self.base_url = "https://api.openai.com/v1/chat/completions"
+            self.provider = "openai"
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+
+    def _prepare_headers(self) -> Dict[str, str]:
+        if self.provider == "anthropic":
+            return {
+                "x-api-key": self.api_key,
+                "anthropic-version": "2023-06-01",
+                "content-type": "application/json",
+            }
+        return {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+    def _prepare_payload(
+        self,
+        messages: List[Dict[str, str]],
+        stream: bool = False,
+    ) -> dict:
+        if self.provider == "anthropic":
+            # Anthropic splits the system prompt out of the messages array.
+            system_msg = next(
+                (m["content"] for m in messages if m["role"] == "system"), ""
+            )
+            user_msgs = [m for m in messages if m["role"] != "system"]
+            return {
+                "model": self.model_id,
+                "system": system_msg,
+                "messages": user_msgs,
+                "max_tokens": 4096,
+                "stream": stream,
+            }
+        return {
+            "model": self.model_id,
+            "messages": messages,
+            "stream": stream,
+        }
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def send_request(self, messages: List[Dict[str, str]]) -> str:
+        """Send a non-streaming chat completion request."""
+        payload = self._prepare_payload(messages, stream=False)
+        headers = self._prepare_headers()
+
+        with httpx.Client(timeout=_DEFAULT_TIMEOUT) as client:
+            response = client.post(self.base_url, json=payload, headers=headers)
+            response.raise_for_status()
+            data = response.json()
+
+        if self.provider == "anthropic":
+            return data["content"][0]["text"]
+        return data["choices"][0]["message"]["content"]
+
+    def stream_request(
+        self, messages: List[Dict[str, str]]
+    ) -> Generator[str, None, None]:
+        """Yield text deltas from an SSE streaming response."""
+        payload = self._prepare_payload(messages, stream=True)
+        headers = self._prepare_headers()
+
+        with httpx.Client(timeout=_DEFAULT_TIMEOUT) as client:
+            with client.stream(
+                "POST", self.base_url, json=payload, headers=headers
+            ) as response:
+                response.raise_for_status()
+                for line in response.iter_lines():
+                    if not line or not line.startswith("data:"):
+                        continue
+                    data_str = line[len("data:"):].strip()
+                    if data_str == "[DONE]":
+                        return
+                    try:
+                        chunk = json.loads(data_str)
+                        delta = chunk["choices"][0].get("delta", {})
+                        content = delta.get("content")
+                        if content:
+                            yield content
+                    except (json.JSONDecodeError, KeyError, IndexError):
+                        # Malformed chunk — skip gracefully.
+                        continue

infinite_context_gateway-0.1.0/infinite_context/cloud/compressor.py

@@ -0,0 +1,44 @@
+"""
+Headroom semantic compression wrapper.
+
+Compresses massive context payloads locally before sending them to a cloud
+LLM, reducing token counts and API costs.
+"""
+
+__all__ = ["HeadroomCompressor"]
+
+
+class HeadroomCompressor:
+    """Wraps the ``headroom-ai`` library with a graceful fallback."""
+
+    def __init__(self, target_ratio: float = 0.8):
+        self.target_ratio = target_ratio
+        self._compress_fn = None
+
+        try:
+            from headroom import compress
+            self._compress_fn = compress
+        except ImportError:
+            print(
+                "Warning: headroom-ai not installed. "
+                "Using fallback semantic compression simulator."
+            )
+
+    def compress(self, context: str) -> str:
+        """Compress *context* using Headroom's algorithms.
+
+        Falls back to a head + tail truncation when the library is absent.
+        """
+        if self._compress_fn is not None:
+            return self._compress_fn(context)
+
+        # Fallback simulator: keep the first 1000 chars and the last 2000
+        # chars (where needles usually are).
+        if len(context) < 3000:
+            return context
+
+        return (
+            context[:1000]
+            + "\n\n...[CONTENT COMPRESSED BY HEADROOM SIMULATOR]...\n\n"
+            + context[-2000:]
+        )

infinite_context_gateway-0.1.0/infinite_context/gateway.py

@@ -0,0 +1,145 @@
+"""
+Context Gateway — unified interface for the cloud and local engines.
+
+Users interact exclusively through :class:`ContextGateway`, selecting either
+``engine="cloud"`` (Headroom compression + API) or ``engine="local"``
+(In-Place Test-Time Training on a local GPU).
+"""
+
+from typing import Any, Generator, List, Optional
+
+__all__ = ["ContextGateway"]
+
+
+class ContextGateway:
+    """Hybrid gateway that routes queries through either a cloud API or a
+    local TTT engine depending on configuration.
+
+    Args:
+        engine: ``"cloud"`` for Headroom API compression, or ``"local"``
+                for In-Place TTT.
+        model_id: Target model (e.g. ``"claude-3-5-sonnet-20240620"`` or a
+                  local HuggingFace repo id).
+        api_key: API key for cloud routing (if applicable).
+        base_url: Custom base URL for OpenAI-compatible providers like Groq.
+        compression_ratio: Target compression ratio for Headroom (cloud only).
+        skip_compression: If *True*, skips Headroom compression (useful for
+                          testing alternative APIs directly).
+        **kwargs: Additional parameters forwarded to :class:`LocalEngine`
+                  (e.g. ``device``, ``load_in_4bit``).
+    """
+
+    def __init__(
+        self,
+        engine: str = "cloud",
+        model_id: str = "claude-3-5-sonnet-20240620",
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        compression_ratio: float = 0.8,
+        skip_compression: bool = False,
+        **kwargs: Any,
+    ):
+        if engine not in ("cloud", "local"):
+            raise ValueError("Engine must be either 'cloud' or 'local'.")
+
+        self.engine = engine
+        self.model_id = model_id
+
+        if self.engine == "cloud":
+            from .cloud.compressor import HeadroomCompressor
+            from .cloud.client import CloudClient
+
+            self.compressor = (
+                HeadroomCompressor(target_ratio=compression_ratio)
+                if not skip_compression
+                else None
+            )
+            self.client = CloudClient(
+                api_key=api_key, model_id=self.model_id, base_url=base_url,
+            )
+        else:  # local
+            from .local.engine import LocalEngine
+            self.local_engine = LocalEngine(model_id=self.model_id, **kwargs)
+
+    # ------------------------------------------------------------------
+    # Chat
+    # ------------------------------------------------------------------
+
+    def chat(
+        self,
+        question: str,
+        context: str,
+        history: Optional[List[dict]] = None,
+        max_history: int = 4,
+        keep_state: bool = False,
+    ) -> str:
+        """Process a massive context payload and answer the question.
+
+        Includes a rolling-window history to keep token costs low.
+        """
+        if self.engine == "cloud":
+            compressed = (
+                self.compressor.compress(context)
+                if self.compressor is not None
+                else context
+            )
+
+            messages: list[dict] = [{"role": "system", "content": compressed}]
+            if history:
+                messages.extend(history[-max_history:])
+            messages.append({"role": "user", "content": question})
+
+            return self.client.send_request(messages)
+
+        # local
+        return self.local_engine.generate(
+            question, context, keep_state=keep_state,
+        )
+
+    def stream_chat(
+        self,
+        question: str,
+        context: str,
+        history: Optional[List[dict]] = None,
+        max_history: int = 4,
+    ) -> Generator[str, None, None]:
+        """Yield response tokens as a stream (cloud engine only)."""
+        if self.engine == "cloud":
+            compressed = (
+                self.compressor.compress(context)
+                if self.compressor is not None
+                else context
+            )
+
+            messages: list[dict] = [{"role": "system", "content": compressed}]
+            if history:
+                messages.extend(history[-max_history:])
+            messages.append({"role": "user", "content": question})
+
+            yield from self.client.stream_request(messages)
+        else:
+            raise NotImplementedError(
+                "Streaming is not supported by the local TTT engine."
+            )
+
+    # ------------------------------------------------------------------
+    # Persistence (local engine only)
+    # ------------------------------------------------------------------
+
+    def save_state(self, path: str) -> None:
+        """Save trained Fast Weights to disk for later resumption."""
+        if self.engine == "local":
+            self.local_engine.save_state(path)
+        else:
+            raise NotImplementedError(
+                "Cloud engine does not support local state saving."
+            )
+
+    def load_state(self, path: str) -> None:
+        """Reload previously-saved Fast Weights from disk."""
+        if self.engine == "local":
+            self.local_engine.load_state(path)
+        else:
+            raise NotImplementedError(
+                "Cloud engine does not support local state loading."
+            )

infinite_context_gateway-0.1.0/infinite_context/local/__init__.py

@@ -0,0 +1,5 @@
+"""Local engine components — In-Place Test-Time Training (TTT) pipeline."""
+
+from .engine import LocalEngine
+
+__all__ = ["LocalEngine"]

infinite_context_gateway-0.1.0/infinite_context/local/chunk_manager.py

@@ -0,0 +1,41 @@
+"""
+Token chunk manager for the local TTT engine.
+
+Splits massive token sequences into fixed-size windows so they fit inside
+the model's KV-Cache limits during the Test-Time Training reading phase.
+"""
+
+from typing import List
+
+__all__ = ["ChunkManager"]
+
+
+class ChunkManager:
+    """Splits token lists into overlapping fixed-size chunks.
+
+    Args:
+        chunk_size: Number of tokens per chunk.
+        overlap: Number of overlapping tokens between consecutive chunks
+                 (helps preserve context continuity).
+    """
+
+    def __init__(self, chunk_size: int = 1024, overlap: int = 0):
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+
+    def chunk_tokens(self, tokens: List[int]) -> List[List[int]]:
+        """Split *tokens* into smaller manageable chunks."""
+        if not tokens:
+            return []
+
+        step = self.chunk_size - self.overlap
+
+        # Ensure we always move forward.
+        if step <= 0:
+            raise ValueError("Overlap cannot be greater than or equal to chunk_size.")
+
+        chunks = []
+        for i in range(0, len(tokens), step):
+            chunks.append(tokens[i : i + self.chunk_size])
+
+        return chunks

infinite_context_gateway-0.1.0/infinite_context/local/engine.py

@@ -0,0 +1,195 @@
+"""
+Local Test-Time Training (TTT) engine.
+
+Orchestrates the full pipeline: model loading → chunking → LoRA injection →
+reading phase (training) → answering phase (inference) → state management.
+"""
+
+import os
+import torch
+from typing import Any
+
+from .chunk_manager import ChunkManager
+from .ttt_module import InPlaceTTT
+
+__all__ = ["LocalEngine"]
+
+# Lazy import guard — resolved once at first use, not at import time.
+_transformers_available: bool | None = None
+
+
+def _ensure_transformers():
+    """Import transformers on first use and raise a clear error if missing."""
+    global _transformers_available
+    if _transformers_available is None:
+        try:
+            import transformers  # noqa: F401
+            _transformers_available = True
+        except ImportError:
+            _transformers_available = False
+    if not _transformers_available:
+        raise ImportError(
+            "Transformers is required for the local engine. "
+            "Install via:  pip install infinite_context[local]"
+        )
+
+
+class LocalEngine:
+    """Orchestrates the Local Test-Time Training (TTT) engine.
+
+    Loads the model in constrained memory (4-bit) and manages the Reading
+    and Answering phases.
+    """
+
+    def __init__(
+        self,
+        model_id: str,
+        device: str = "cuda",
+        load_in_4bit: bool = True,
+        **kwargs: Any,
+    ):
+        _ensure_transformers()
+        
+        self.device = device if torch.cuda.is_available() else "cpu"
+        self.model_id = model_id
+
+        print(f"Loading base model {model_id} via Transformers into {self.device}...")
+        
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        
+        quantization_kwargs: dict = {}
+        if load_in_4bit and self.device == "cuda":
+            try:
+                from transformers import BitsAndBytesConfig
+                quantization_kwargs["quantization_config"] = BitsAndBytesConfig(
+                    load_in_4bit=True,
+                    bnb_4bit_compute_dtype=torch.float16,
+                    bnb_4bit_use_double_quant=True,
+                )
+            except ImportError:
+                pass
+
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_id,
+            device_map="auto" if self.device == "cuda" else None,
+            **quantization_kwargs,
+        )
+
+        self.chunker = ChunkManager(chunk_size=1024, overlap=128)
+
+    # ------------------------------------------------------------------
+    # Core pipeline
+    # ------------------------------------------------------------------
+
+    def generate(
+        self,
+        question: str,
+        context: str,
+        epochs_per_chunk: int = 15,
+        target_loss: float = 0.05,
+        keep_state: bool = False,
+    ) -> str:
+        """Run the full In-Place TTT pipeline.
+
+        1. Break massive context into chunks.
+        2. Inject Fast Weights (LoRA) into the model.
+        3. Reading Phase — train the Fast Weights on the chunks.
+        4. Answering Phase — generate the response to the question.
+        5. State Management — optionally retain or destroy Fast Weights.
+        """
+        # 1. Chunking
+        tokens = self.tokenizer.encode(context, add_special_tokens=False)
+        chunks = self.chunker.chunk_tokens(tokens)
+
+        print(f"Bypassing KV-Cache: Splitting {len(tokens)} tokens into {len(chunks)} chunks.")
+
+        # 2. Attach TTT Fast Weights (skip if a checkpoint was loaded)
+        is_peft = hasattr(self.model, "peft_config")
+        ttt: InPlaceTTT | None = None
+        if not is_peft:
+            print("Injecting TTT Fast Weights into down_proj layers...")
+            ttt = InPlaceTTT(self.model)
+        else:
+            print("Model already has Fast Weights loaded from disk. Reusing them...")
+
+        # 3. The 'Reading' Phase
+        if chunks:
+            print("Starting Reading Phase (Test-Time Training)...")
+            if ttt is None:
+                print("Warning: Cannot train loaded Fast Weights in-place. Skipping training.")
+            else:
+                for i, chunk in enumerate(chunks):
+                    # Hoist tensor creation outside the epoch loop.
+                    input_ids = torch.tensor([chunk], device=self.device)
+                    for epoch in range(epochs_per_chunk):
+                        loss = ttt.train_on_chunk(input_ids)
+                        if loss < target_loss:
+                            break
+                    print(f"  Processed Chunk {i + 1}/{len(chunks)} - Epochs: {epoch + 1} - Loss: {loss:.4f}")
+
+        # KV Cache is implicitly cleared here because we are not passing
+        # past_key_values!
+
+        # 4. The 'Answering' Phase
+        print("Starting Answering Phase...")
+        self.model.eval()
+
+        prompt = f"Based on the codebase I just showed you, answer this question:\n{question}\nAnswer:"
+        prompt_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
+
+        device_type = self.device if self.device in ("cuda", "cpu") else "cpu"
+        with torch.no_grad():
+            output_ids = self.model.generate(
+                prompt_ids,
+                max_new_tokens=75,
+                do_sample=False,
+                repetition_penalty=1.1,
+            )
+
+        # Extract just the newly generated tokens.
+        response_tokens = output_ids[0][prompt_ids.shape[1]:]
+        response = self.tokenizer.decode(response_tokens, skip_special_tokens=True)
+
+        # 5. State Management
+        if not keep_state:
+            if ttt is not None:
+                print("Resetting model state (deleting Fast Weights)...")
+                ttt.remove_fast_weights()
+            else:
+                print("Cannot reset loaded Fast Weights. Retaining state.")
+        else:
+            print("Keeping model state (Fast Weights retained)...")
+            if ttt is not None:
+                # Update the reference so save_pretrained saves the adapter.
+                self.model = ttt.model
+
+        return response
+
+    # ------------------------------------------------------------------
+    # Persistence
+    # ------------------------------------------------------------------
+
+    def save_state(self, path: str) -> None:
+        """Save the current Fast Weights (LoRA adapter) to disk."""
+        if hasattr(self.model, "peft_config"):
+            self.model.save_pretrained(path)
+            print(f"State saved to {path}")
+        else:
+            raise RuntimeError(
+                "No active Fast Weights to save. "
+                "Did you forget keep_state=True when calling generate()?"
+            )
+
+    def load_state(self, path: str) -> None:
+        """Load Fast Weights (LoRA adapter) from disk."""
+        from peft import PeftModel
+
+        adapter_path = os.path.join(path, "ttt_fast_weights")
+        if not os.path.exists(adapter_path):
+            # Fallback for standard PEFT saves.
+            adapter_path = path
+        self.model = PeftModel.from_pretrained(
+            self.model, adapter_path, is_trainable=True,
+        )
+        print(f"State loaded from {adapter_path}")

infinite_context_gateway-0.1.0/infinite_context/local/ttt_module.py

@@ -0,0 +1,133 @@
+"""
+In-Place Test-Time Training (TTT) module.
+
+Attaches transient LoRA fast-weight adapters to an LLM's MLP layers and
+trains them on incoming context chunks so the model "memorises" the document
+into its weights at inference time.
+"""
+
+import torch
+import torch.nn as nn
+from typing import List, Optional
+
+__all__ = ["InPlaceTTT"]
+
+# Lazy guard — resolved once at first use, not at import time.
+_peft_available: Optional[bool] = None
+
+
+def _ensure_peft():
+    """Import peft on first use and raise a clear error if missing."""
+    global _peft_available
+    if _peft_available is None:
+        try:
+            import peft  # noqa: F401
+            _peft_available = True
+        except ImportError:
+            _peft_available = False
+    if not _peft_available:
+        raise ImportError(
+            "PEFT is required for the local engine. "
+            "Install via:  pip install infinite_context[local]"
+        )
+
+
+class InPlaceTTT:
+    """Manages LoRA adapter injection, test-time training, and teardown."""
+
+    def __init__(
+        self,
+        model: nn.Module,
+        target_modules: Optional[List[str]] = None,
+        lr: float = 5e-3,
+        adapter_name: str = "ttt_fast_weights",
+    ):
+        """
+        Args:
+            model: The Hugging Face PreTrainedModel.
+            target_modules: MLP layer names to attach fast weights to.
+                            Defaults to ``["down_proj"]``.
+            lr: Learning rate for the test-time training update.
+            adapter_name: Name of the PEFT adapter used for the fast weights.
+        """
+        _ensure_peft()
+
+        self.model = model
+        self.target_modules = target_modules or ["down_proj"]
+        self.lr = lr
+        self.adapter_name = adapter_name
+        self.optimizer: Optional[torch.optim.Optimizer] = None
+
+        self._setup_lora()
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+
+    def _setup_lora(self) -> None:
+        """Inject Fast Weights (LoRA) into the base model.
+
+        By setting ``r=8`` and targeting ``down_proj``, we turn the static
+        MLP into an active, online memory bank.
+        """
+        from peft import LoraConfig, get_peft_model
+        config = LoraConfig(
+            r=8,
+            lora_alpha=16,
+            target_modules=self.target_modules,
+            lora_dropout=0.0,
+            bias="none",
+            task_type="CAUSAL_LM",
+            use_dora=True,
+        )
+        self.model = get_peft_model(self.model, config, adapter_name=self.adapter_name)
+
+        # Train ONLY the LoRA parameters — not the frozen base weights.
+        trainable_params = [p for p in self.model.parameters() if p.requires_grad]
+        self.optimizer = torch.optim.AdamW(trainable_params, lr=self.lr)
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def train_on_chunk(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> float:
+        """Run one gradient step on *input_ids* (the 'Reading Phase').
+
+        Uses the standard causal next-token prediction loss which achieves
+        the same goal as the paper's Conv1D LM-Aligned objective: baking
+        the chunk's knowledge into the ``down_proj`` weights.
+        """
+        self.model.train()
+        self.optimizer.zero_grad(set_to_none=True)
+
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        # Mixed-precision forward pass (FP16 on CUDA, no-op on CPU).
+        device_type = input_ids.device.type
+        with torch.amp.autocast(device_type, enabled=(device_type == "cuda")):
+            outputs = self.model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=input_ids,
+            )
+            loss = outputs.loss
+
+        # Backward + update ONLY the LoRA adapters.
+        loss.backward()
+        self.optimizer.step()
+
+        return loss.item()
+
+    def remove_fast_weights(self) -> None:
+        """Delete the LoRA adapter ('State Destruction').
+
+        The model instantly resets to its base pre-trained state, ready for
+        the next request.
+        """
+        self.model.eval()
+        self.model.delete_adapter(self.adapter_name)

infinite_context_gateway-0.1.0/infinite_context_gateway.egg-info/PKG-INFO

@@ -0,0 +1,108 @@
+Metadata-Version: 2.4
+Name: infinite-context-gateway
+Version: 0.1.0
+Summary: A hybrid AI context gateway combining Headroom cloud compression and In-Place Local Test-Time Training (TTT).
+Author-email: Dev Team <dev@example.com>
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: httpx
+Provides-Extra: cloud
+Requires-Dist: headroom-ai; extra == "cloud"
+Provides-Extra: local
+Requires-Dist: torch; extra == "local"
+Requires-Dist: transformers; extra == "local"
+Requires-Dist: peft; extra == "local"
+Requires-Dist: accelerate; extra == "local"
+Requires-Dist: bitsandbytes; extra == "local"
+Provides-Extra: all
+Requires-Dist: headroom-ai; extra == "all"
+Requires-Dist: torch; extra == "all"
+Requires-Dist: transformers; extra == "all"
+Requires-Dist: peft; extra == "all"
+Requires-Dist: accelerate; extra == "all"
+Requires-Dist: bitsandbytes; extra == "all"
+
+# Infinite Context
+
+A hybrid AI context gateway combining **Headroom cloud compression** and **In-Place Local Test-Time Training (TTT)**.
+
+## Installation
+
+To install the base gateway:
+```bash
+pip install infinite_context
+```
+
+To install with Cloud Phase 1 dependencies (Headroom API):
+```bash
+pip install infinite_context[cloud]
+```
+
+To install with Local Phase 2 dependencies (PyTorch, Transformers, PEFT):
+```bash
+pip install infinite_context[local]
+```
+
+## Usage
+
+### Phase 1: Cloud Compression
+Uses the Headroom API to semantically compress a massive context and send it to cloud providers (OpenAI, Anthropic) while keeping costs low.
+
+```python
+from infinite_context import ContextGateway
+
+gateway = ContextGateway(
+    engine="cloud",
+    model_id="claude-3-5-sonnet-20240620",
+    api_key="your_anthropic_api_key",
+    compression_ratio=0.8
+)
+
+# You can pass in conversational history (rolling window)
+history = [
+    {"role": "user", "content": "What is the codebase about?"},
+    {"role": "assistant", "content": "It is a Python application..."}
+]
+
+response = gateway.chat("How does the failover protocol work?", massive_context, history=history)
+print(response)
+```
+
+### Phase 2: Local Test-Time Training (TTT)
+Bypasses the KV-Cache entirely by injecting a PEFT LoRA adapter and baking the context directly into the model's neural weights on your local GPU. Includes Early Stopping latency optimizations and Generation repetition penalties.
+
+```python
+from infinite_context import ContextGateway
+
+gateway = ContextGateway(
+    engine="local",
+    model_id="Qwen/Qwen2.5-0.5B-Instruct",
+    load_in_4bit=True
+)
+
+response = gateway.chat("What is the failover protocol?", massive_context)
+print(response)
+```
+
+### Checkpoint Persistence
+Save trained Fast Weights to disk and resume later without re-reading the document:
+
+```python
+from infinite_context import ContextGateway
+
+# Train and keep state
+gateway = ContextGateway(engine="local", model_id="Qwen/Qwen2.5-0.5B-Instruct")
+gateway.chat("Summarise the document.", massive_context, keep_state=True)
+gateway.save_state("./my_checkpoint")
+
+del gateway  # Free all GPU memory
+
+# Resume later — zero re-training
+gateway = ContextGateway(engine="local", model_id="Qwen/Qwen2.5-0.5B-Instruct")
+gateway.load_state("./my_checkpoint")
+response = gateway.chat("What was the protocol?", context="")
+print(response)
+```

infinite_context_gateway-0.1.0/infinite_context_gateway.egg-info/SOURCES.txt

@@ -0,0 +1,16 @@
+README.md
+pyproject.toml
+infinite_context/__init__.py
+infinite_context/gateway.py
+infinite_context/cloud/__init__.py
+infinite_context/cloud/client.py
+infinite_context/cloud/compressor.py
+infinite_context/local/__init__.py
+infinite_context/local/chunk_manager.py
+infinite_context/local/engine.py
+infinite_context/local/ttt_module.py
+infinite_context_gateway.egg-info/PKG-INFO
+infinite_context_gateway.egg-info/SOURCES.txt
+infinite_context_gateway.egg-info/dependency_links.txt
+infinite_context_gateway.egg-info/requires.txt
+infinite_context_gateway.egg-info/top_level.txt

infinite_context_gateway-0.1.0/infinite_context_gateway.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

infinite_context_gateway-0.1.0/infinite_context_gateway.egg-info/requires.txt

@@ -0,0 +1,19 @@
+httpx
+
+[all]
+headroom-ai
+torch
+transformers
+peft
+accelerate
+bitsandbytes
+
+[cloud]
+headroom-ai
+
+[local]
+torch
+transformers
+peft
+accelerate
+bitsandbytes

infinite_context_gateway-0.1.0/infinite_context_gateway.egg-info/top_level.txt

@@ -0,0 +1 @@
+infinite_context

infinite_context_gateway-0.1.0/pyproject.toml

@@ -0,0 +1,44 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "infinite-context-gateway"
+version = "0.1.0"
+authors = [
+  { name="Dev Team", email="dev@example.com" },
+]
+description = "A hybrid AI context gateway combining Headroom cloud compression and In-Place Local Test-Time Training (TTT)."
+readme = "README.md"
+requires-python = ">=3.9"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+    "httpx",
+]
+
+[project.optional-dependencies]
+cloud = [
+    "headroom-ai",
+]
+local = [
+    "torch",
+    "transformers",
+    "peft",
+    "accelerate",
+    "bitsandbytes",
+]
+all = [
+    "headroom-ai",
+    "torch",
+    "transformers",
+    "peft",
+    "accelerate",
+    "bitsandbytes",
+]
+
+[tool.setuptools.packages.find]
+include = ["infinite_context*"]

infinite_context_gateway-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+