Add Kompress: ModernBERT token compressor replacing LLMLingua-2

Adds kompress_compressor.py — a self-contained ModernBERT-based token
compressor that auto-downloads from chopratejas/kompress-base on HuggingFace.
Trained on 330K structured tool outputs (JSON, diffs, logs, code, SQL,
agentic traces), achieving 82% entity preservation vs LLMLingua-2's 36%.

Changes:
- New: kompress_compressor.py — dual-head ModernBERT (token + span CNN)
with HuggingFace auto-download, no extra pip install needed
- ContentRouter: Kompress is primary ML compressor, LLMLingua-2 is fallback
- fallback_strategy changed from PASSTHROUGH to KOMPRESS — unknown/mixed
content now gets compressed instead of ignored
- No hardcoded compression ratios — model decides per-token importance,
optional target_ratio only when user explicitly sets it via API
- Version bump: 0.3.8 → 0.4.0

headroom/__init__.py

@@ -142,7 +142,7 @@ from .transforms import (
     TransformPipeline,
 )
 
-__version__ = "0.3.7"
+__version__ = "0.4.0"
 
 __all__ = [
     # Main client

headroom/transforms/content_router.py

@@ -246,6 +246,7 @@ class CompressionStrategy(Enum):
     SMART_CRUSHER = "smart_crusher"
     SEARCH = "search"
     LOG = "log"
+    KOMPRESS = "kompress"
     LLMLINGUA = "llmlingua"
     TEXT = "text"
     DIFF = "diff"
@@ -371,6 +372,7 @@ class ContentRouterConfig:
 
     # Enable/disable specific compressors
     enable_code_aware: bool = True
+    enable_kompress: bool = True  # Kompress: ModernBERT token compressor (preferred over LLMLingua)
     enable_llmlingua: bool = True
     enable_smart_crusher: bool = True
     enable_search_compressor: bool = True
@@ -383,8 +385,8 @@ class ContentRouterConfig:
     mixed_content_threshold: int = 2  # Min types to consider mixed
     min_section_tokens: int = 20  # Min tokens to compress a section
 
-    # Fallback
-    fallback_strategy: CompressionStrategy = CompressionStrategy.PASSTHROUGH
+    # Fallback: Kompress handles unknown/mixed content instead of passing through
+    fallback_strategy: CompressionStrategy = CompressionStrategy.KOMPRESS
 
     # Protection: Don't compress content that's likely the subject of analysis
     skip_user_messages: bool = True  # User messages contain what they want analyzed
@@ -633,6 +635,7 @@ class ContentRouter(Transform):
         self._log_compressor: Any = None
         self._diff_compressor: Any = None
         self._html_extractor: Any = None
+        self._kompress: Any = None
         self._llmlingua: Any = None
         self._text_compressor: Any = None
         self._image_optimizer: Any = None
@@ -989,13 +992,16 @@ class ContentRouter(Transform):
                         # Estimate tokens from extracted text (simple word count)
                         compressed_tokens = len(compressed.split()) if compressed else 0
 
+            elif strategy == CompressionStrategy.KOMPRESS:
+                compressed, compressed_tokens = self._try_ml_compressor(content, context, question)
+
             elif strategy == CompressionStrategy.LLMLINGUA:
-                compressed, compressed_tokens = self._try_llmlingua(content, context, question)
+                compressed, compressed_tokens = self._try_ml_compressor(content, context, question)
 
             elif strategy == CompressionStrategy.TEXT:
-                # Prefer LLMLingua for text if available (ML-based, better compression)
-                # Falls back to heuristic TextCompressor if LLMLingua unavailable
-                compressed, compressed_tokens = self._try_llmlingua(content, context, question)
+                # Prefer ML compressor (Kompress > LLMLingua) for text
+                # Falls back to heuristic TextCompressor if neither available
+                compressed, compressed_tokens = self._try_ml_compressor(content, context, question)
 
         except Exception as e:
             logger.warning("Compression with %s failed: %s", strategy.value, e)
@@ -1016,10 +1022,13 @@ class ContentRouter(Transform):
         # Fallback: return unchanged
         return content, original_tokens
 
-    def _try_llmlingua(
+    def _try_ml_compressor(
         self, content: str, context: str, question: str | None = None
     ) -> tuple[str, int]:
-        """Try LLMLingua compression with fallback.
+        """ML-based compression: Kompress (primary), LLMLingua (fallback only).
+
+        Kompress (ModernBERT, trained on 330K structured tool outputs)
+        auto-downloads from HuggingFace on first use. No heuristic fallback.
 
         Args:
             content: Content to compress.
@@ -1029,6 +1038,17 @@ class ContentRouter(Transform):
         Returns:
             Tuple of (compressed, token_count).
         """
+        # Primary: Kompress — downloads from chopratejas/kompress-base on first use
+        if self.config.enable_kompress:
+            compressor = self._get_kompress()
+            if compressor:
+                try:
+                    result = compressor.compress(content, context=context, question=question)
+                    return result.compressed, result.compressed_tokens
+                except Exception as e:
+                    logger.warning("Kompress failed: %s", e)
+
+        # Fallback: LLMLingua (only if Kompress not installed)
         if self.config.enable_llmlingua:
             compressor = self._get_llmlingua()
             if compressor:
@@ -1036,16 +1056,13 @@ class ContentRouter(Transform):
                     result = compressor.compress(content, context=context, question=question)
                     return result.compressed, result.compressed_tokens
                 except Exception as e:
-                    logger.debug("LLMLingua failed: %s", e)
-
-        # Fallback to text compressor
-        compressor = self._get_text_compressor()
-        if compressor:
-            result = compressor.compress(content, context=context)
-            return result.compressed, result.compressed_line_count
+                    logger.warning("LLMLingua failed: %s", e)
 
         return content, len(content.split())
 
+    # Backwards compatibility
+    _try_llmlingua = _try_ml_compressor
+
     def _strategy_from_detection_type(self, content_type: ContentType) -> CompressionStrategy:
         """Get strategy from ContentType enum."""
         mapping = {
@@ -1069,6 +1086,7 @@ class ContentRouter(Transform):
             CompressionStrategy.DIFF: ContentType.GIT_DIFF,
             CompressionStrategy.HTML: ContentType.HTML,
             CompressionStrategy.TEXT: ContentType.PLAIN_TEXT,
+            CompressionStrategy.KOMPRESS: ContentType.PLAIN_TEXT,
             CompressionStrategy.LLMLINGUA: ContentType.PLAIN_TEXT,
             CompressionStrategy.PASSTHROUGH: ContentType.PLAIN_TEXT,
         }
@@ -1154,13 +1172,19 @@ class ContentRouter(Transform):
     def eager_load_compressors(self) -> None:
         """Pre-load compressors at startup to avoid first-request latency.
 
-        Call this during proxy startup to load LLMLingua model (~5s)
+        Call this during proxy startup to load models (~5s)
         before any requests arrive.
         """
+        # Prefer Kompress (faster, smaller, better on structured data)
+        if self.config.enable_kompress:
+            compressor = self._get_kompress()
+            if compressor:
+                logger.info("Kompress model pre-loaded at startup")
+                return  # No need to also load LLMLingua
+
         if self.config.enable_llmlingua:
             compressor = self._get_llmlingua()
             if compressor:
-                # Trigger the underlying model load by accessing it
                 try:
                     from .llmlingua_compressor import _get_llmlingua_compressor
 
@@ -1170,8 +1194,20 @@ class ContentRouter(Transform):
                 except Exception as e:
                     logger.warning("Failed to pre-load LLMLingua model: %s", e)
 
+    def _get_kompress(self) -> Any:
+        """Get KompressCompressor (lazy load). Downloads from HuggingFace on first use."""
+        if self._kompress is None:
+            try:
+                from .kompress_compressor import KompressCompressor, is_kompress_available
+
+                if is_kompress_available():
+                    self._kompress = KompressCompressor()
+            except ImportError:
+                logger.debug("Kompress dependencies not available")
+        return self._kompress
+
     def _get_llmlingua(self) -> Any:
-        """Get LLMLinguaCompressor (lazy load)."""
+        """Get LLMLinguaCompressor (lazy load). Fallback if Kompress unavailable."""
         if self._llmlingua is None:
             try:
                 from .llmlingua_compressor import (

headroom/transforms/kompress_compressor.py

@@ -0,0 +1,363 @@
+"""Kompress: ModernBERT token compressor for structured tool outputs.
+
+Drop-in replacement for LLMLingua-2. Auto-downloads the model from
+HuggingFace (chopratejas/kompress-base) on first use.
+
+No extra pip install needed — uses transformers + safetensors
+which are already Headroom dependencies.
+
+Usage:
+    >>> from headroom.transforms.kompress_compressor import KompressCompressor
+    >>> compressor = KompressCompressor()
+    >>> result = compressor.compress(long_tool_output)
+    >>> print(result.compressed)
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+from transformers import AutoModel, AutoTokenizer
+
+from ..config import TransformResult
+from ..tokenizer import Tokenizer
+from .base import Transform
+
+logger = logging.getLogger(__name__)
+
+# HuggingFace model ID
+HF_MODEL_ID = "chopratejas/kompress-base"
+
+# Lazy singleton
+_kompress_model = None
+_kompress_tokenizer = None
+_kompress_lock = threading.Lock()
+
+
+# ── Model Architecture (must match training) ──────────────────────────
+
+
+class HeadroomCompressorModel(nn.Module):
+    """Dual-head ModernBERT: token classification + span importance CNN."""
+
+    def __init__(self, model_name: str = "answerdotai/ModernBERT-base"):
+        super().__init__()
+        self.encoder = AutoModel.from_pretrained(model_name, attn_implementation="eager")
+        hidden_size = self.encoder.config.hidden_size  # 768
+
+        # Head 1: Token keep/discard
+        self.token_dropout = nn.Dropout(0.1)
+        self.token_head = nn.Linear(hidden_size, 2)
+
+        # Head 2: Span importance (1D CNN)
+        self.span_conv = nn.Sequential(
+            nn.Conv1d(hidden_size, 256, kernel_size=5, padding=2),
+            nn.GELU(),
+            nn.Conv1d(256, 1, kernel_size=3, padding=1),
+            nn.Sigmoid(),
+        )
+
+    def get_scores(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        """Get per-token compression scores. Higher = more important."""
+        with torch.no_grad():
+            hidden = self.encoder(input_ids, attention_mask=attention_mask).last_hidden_state
+
+            token_probs = torch.softmax(self.token_head(hidden), dim=-1)[:, :, 1]
+
+            span_scores = self.span_conv(hidden.transpose(1, 2)).squeeze(1)
+
+            return token_probs * (0.5 + 0.5 * span_scores)
+
+
+# ── Model Loading ─────────────────────────────────────────────────────
+
+
+def _load_kompress(device: str = "auto") -> tuple[HeadroomCompressorModel, Any]:
+    """Download from HuggingFace and load the Kompress model."""
+    global _kompress_model, _kompress_tokenizer
+
+    with _kompress_lock:
+        if _kompress_model is not None:
+            return _kompress_model, _kompress_tokenizer
+
+        from huggingface_hub import hf_hub_download
+
+        logger.info("Downloading Kompress model from %s ...", HF_MODEL_ID)
+
+        # Download model weights
+        weights_path = hf_hub_download(HF_MODEL_ID, "model.safetensors")
+
+        # Load architecture
+        model = HeadroomCompressorModel()
+
+        # Load trained weights
+        from safetensors.torch import load_file
+
+        state_dict = load_file(weights_path)
+        model.load_state_dict(state_dict, strict=False)
+
+        # Resolve device
+        if device == "auto":
+            if torch.cuda.is_available():
+                device = "cuda"
+            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                device = "mps"
+            else:
+                device = "cpu"
+
+        model.to(device)
+        model.eval()
+        logger.info("Kompress model loaded on %s (%s)", device, HF_MODEL_ID)
+
+        tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
+
+        _kompress_model = model
+        _kompress_tokenizer = tokenizer
+        return model, tokenizer
+
+
+def is_kompress_available() -> bool:
+    """Check if Kompress dependencies are available."""
+    try:
+        import huggingface_hub  # noqa: F401
+        import safetensors  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
+def unload_kompress_model() -> bool:
+    """Unload the Kompress model to free memory."""
+    global _kompress_model, _kompress_tokenizer
+    with _kompress_lock:
+        if _kompress_model is not None:
+            _kompress_model = None
+            _kompress_tokenizer = None
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            return True
+    return False
+
+
+# ── Compressor ────────────────────────────────────────────────────────
+
+
+@dataclass
+class KompressConfig:
+    """Minimal config. The model decides what's important — not us."""
+
+    device: str = "auto"
+    enable_ccr: bool = True
+
+
+@dataclass
+class KompressResult:
+    """Result of Kompress compression."""
+
+    compressed: str
+    original: str
+    original_tokens: int
+    compressed_tokens: int
+    compression_ratio: float
+    cache_key: str | None = None
+    model_used: str = HF_MODEL_ID
+
+    @property
+    def tokens_saved(self) -> int:
+        return max(0, self.original_tokens - self.compressed_tokens)
+
+    @property
+    def savings_percentage(self) -> float:
+        if self.original_tokens == 0:
+            return 0.0
+        return (self.tokens_saved / self.original_tokens) * 100
+
+
+class KompressCompressor(Transform):
+    """Kompress: ModernBERT token compressor for structured tool outputs.
+
+    Auto-downloads chopratejas/kompress-base from HuggingFace on first use.
+    Drop-in replacement for LLMLinguaCompressor with identical interface.
+    """
+
+    name: str = "kompress_compressor"
+
+    def __init__(self, config: KompressConfig | None = None):
+        self.config = config or KompressConfig()
+
+    def compress(
+        self,
+        content: str,
+        context: str = "",
+        content_type: str | None = None,
+        question: str | None = None,
+        target_ratio: float | None = None,
+    ) -> KompressResult:
+        """Compress content using Kompress model.
+
+        Args:
+            content: Text to compress.
+            context: Optional surrounding context (unused by model, kept for interface compat).
+            content_type: Ignored — model decides importance per content type.
+            question: Ignored — kept for LLMLingua interface compat.
+            target_ratio: If None (default), model decides how much to keep using
+                score threshold. If set (e.g. 0.3), forces that keep ratio.
+                The proxy never sets this — only user-facing API does.
+
+        Returns:
+            KompressResult with compressed text.
+        """
+        words = content.split()
+        n_words = len(words)
+
+        if n_words < 10:
+            return self._passthrough(content, n_words)
+
+        try:
+            model, tokenizer = _load_kompress(self.config.device)
+
+            # Tokenize
+            encoding = tokenizer(
+                words,
+                is_split_into_words=True,
+                truncation=True,
+                max_length=8192,
+                padding=True,
+                return_tensors="pt",
+            )
+
+            device = next(model.parameters()).device
+            input_ids = encoding["input_ids"].to(device)
+            attention_mask = encoding["attention_mask"].to(device)
+
+            # Get per-token importance scores from dual-head model
+            scores = model.get_scores(input_ids, attention_mask)[0].cpu()
+
+            # Map subword scores to word-level (max pooling)
+            word_ids = encoding.word_ids(batch_index=0)
+            word_scores: dict[int, float] = {}
+            for idx, wid in enumerate(word_ids):
+                if wid is None:
+                    continue
+                s = scores[idx].item()
+                if wid not in word_scores or s > word_scores[wid]:
+                    word_scores[wid] = s
+
+            if not word_scores:
+                return self._passthrough(content, n_words)
+
+            # Token selection
+            if target_ratio is not None:
+                # User explicitly asked for a specific ratio — use top-k
+                sorted_wids = sorted(word_scores, key=lambda w: word_scores[w], reverse=True)
+                num_keep = max(1, int(len(sorted_wids) * target_ratio))
+                kept_ids = set(sorted_wids[:num_keep])
+            else:
+                # Model decides. Trained with binary labels — score > 0.5 = keep.
+                # Dense content → most tokens score high → keeps more.
+                # Boilerplate → most score low → keeps less. That's correct.
+                kept_ids = {wid for wid, score in word_scores.items() if score > 0.5}
+                if not kept_ids:
+                    # Edge case: nothing above threshold — keep the single highest
+                    best = max(word_scores, key=lambda w: word_scores[w])
+                    kept_ids = {best}
+
+            # Reconstruct in original word order
+            compressed_words = [words[w] for w in sorted(kept_ids) if w < n_words]
+            compressed = " ".join(compressed_words)
+            compressed_count = len(compressed_words)
+            ratio = compressed_count / n_words if n_words else 1.0
+
+            result = KompressResult(
+                compressed=compressed,
+                original=content,
+                original_tokens=n_words,
+                compressed_tokens=compressed_count,
+                compression_ratio=ratio,
+            )
+
+            # CCR marker
+            if self.config.enable_ccr and ratio < 0.8:
+                cache_key = self._store_in_ccr(content, compressed, n_words)
+                if cache_key:
+                    result.cache_key = cache_key
+                    result.compressed += (
+                        f"\n[{n_words} items compressed to {compressed_count}."
+                        f" Retrieve more: hash={cache_key}]"
+                    )
+
+            return result
+
+        except Exception as e:
+            logger.warning("Kompress compression failed: %s", e)
+            return self._passthrough(content, n_words)
+
+    def _passthrough(self, content: str, n_words: int) -> KompressResult:
+        return KompressResult(
+            compressed=content,
+            original=content,
+            original_tokens=n_words,
+            compressed_tokens=n_words,
+            compression_ratio=1.0,
+        )
+
+    def apply(
+        self,
+        messages: list[dict[str, Any]],
+        tokenizer: Tokenizer,
+        **kwargs: Any,
+    ) -> TransformResult:
+        """Apply Kompress compression to messages (Transform interface)."""
+        tokens_before = sum(tokenizer.count_text(str(m.get("content", ""))) for m in messages)
+        transformed = []
+        transforms_applied = []
+
+        for message in messages:
+            role = message.get("role", "")
+            content = message.get("content", "")
+
+            if not isinstance(content, str) or len(content.split()) < 10:
+                transformed.append(message)
+                continue
+
+            # Compress tool outputs and long assistant messages
+            # Model decides how much — no hardcoded ratios
+            if role in ("tool", "assistant"):
+                result = self.compress(content)
+                if result.compression_ratio < 0.9:
+                    transformed.append({**message, "content": result.compressed})
+                    transforms_applied.append(f"kompress:{role}:{result.compression_ratio:.2f}")
+                else:
+                    transformed.append(message)
+            else:
+                transformed.append(message)
+
+        tokens_after = sum(tokenizer.count_text(str(m.get("content", ""))) for m in transformed)
+
+        return TransformResult(
+            messages=transformed,
+            tokens_before=tokens_before,
+            tokens_after=tokens_after,
+            transforms_applied=transforms_applied or ["kompress:noop"],
+        )
+
+    def _store_in_ccr(self, original: str, compressed: str, original_tokens: int) -> str | None:
+        try:
+            from ..cache.compression_store import get_compression_store
+
+            store = get_compression_store()
+            return store.store(
+                original,
+                compressed,
+                original_tokens=original_tokens,
+                compressed_tokens=len(compressed.split()),
+                compression_strategy="kompress",
+            )
+        except Exception:
+            return None

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "headroom-ai"
-version = "0.3.7"
+version = "0.4.0"
 description = "The Context Optimization Layer for LLM Applications - Cut costs by 50-90%"
 readme = "README.md"
 license = "Apache-2.0"
@@ -76,7 +76,7 @@ proxy = [
 reports = [
     "jinja2>=3.0.0",
 ]
-# ML-based compression (LLMLingua-2)
+# ML-based compression (LLMLingua-2 — fallback if Kompress fails)
 llmlingua = [
     "llmlingua>=0.2.0",
     "torch>=2.0.0",

tests/test_transforms/test_content_router.py

@@ -155,12 +155,13 @@ class TestContentRouterConfig:
         config = ContentRouterConfig()
 
         assert config.enable_code_aware is True
+        assert config.enable_kompress is True
         assert config.enable_llmlingua is True
         assert config.enable_smart_crusher is True
         assert config.enable_search_compressor is True
         assert config.enable_log_compressor is True
         assert config.min_section_tokens == 20
-        assert config.fallback_strategy == CompressionStrategy.PASSTHROUGH
+        assert config.fallback_strategy == CompressionStrategy.KOMPRESS
 
     def test_custom_values(self):
         """Custom config values are applied."""