"""Kompress: ModernBERT token compressor for structured tool outputs. Auto-downloads the model from HuggingFace (chopratejas/kompress-v2-base) on first use. Requires the [ml] extra: pip install headroom-ai[ml] Usage: >>> from headroom.transforms.kompress_compressor import KompressCompressor >>> compressor = KompressCompressor() >>> result = compressor.compress(long_tool_output) >>> print(result.compressed) """ from __future__ import annotations import contextlib import gc import hashlib import logging import os import re import threading import time from dataclasses import dataclass from typing import Any, Literal from ..config import TransformResult from ..onnx_runtime import ( create_cpu_session_options, hf_hub_download_local_first, trim_process_heap, ) from ..tokenizer import Tokenizer from .base import Transform logger = logging.getLogger(__name__) # Default HuggingFace model ID HF_MODEL_ID = "chopratejas/kompress-v2-base" # Tokens matching this pattern are always kept regardless of model score. # Numbers, ALLCAPS identifiers, dotted paths, unix paths, file extensions, # CLI flags, and CamelCase names carry semantic meaning that agents cannot # reconstruct from context — dropping them degrades reasoning correctness. # Disable with HEADROOM_KOMPRESS_MUST_KEEP=0. _KOMPRESS_MUST_KEEP_RE = re.compile( r"\b0x[0-9A-Fa-f]+\b" # hex addresses/IDs: 0x7fff2038 r"|(? None: """Add semantically fragile words that should never be model-dropped.""" if os.environ.get(_KOMPRESS_MUST_KEEP_ENV, "1") == "0": return for word_idx, word in enumerate(chunk_words): if _KOMPRESS_MUST_KEEP_RE.search(word): kept_ids.add(word_idx + chunk_start) # ONNX artifacts are resolved against the model repo in this order, falling # through on download miss OR session-load failure: # # - kompress-int8-wo.onnx: weight-only int8 (MatMulNBits), 261MB. Evaluated on # the labeled dataset_v2 test split (n=500): f1=0.9130 vs fp32's 0.9128, # must_keep_recall 0.9765 vs 0.9770, keep_rate 0.8097 vs 0.8100, 99.6% # keep-decision agreement — fp32-equivalent at 2.2x less memory. Uses the # com.microsoft MatMulNBits contrib op; older onnxruntime builds without the # 8-bit kernel fail at session load and fall through to fp32. # - kompress-fp32.onnx: lossless reference, 601MB. # - kompress-int8.onnx: v1-era dynamic int8 (kept for custom domain repos). # # An operator can pin an exact file via HEADROOM_KOMPRESS_ONNX_FILENAME. _DEFAULT_ONNX_FILENAMES = ( "onnx/kompress-int8-wo.onnx", "onnx/kompress-fp32.onnx", "onnx/kompress-int8.onnx", ) KOMPRESS_ONNX_INTRA_THREADS_ENV = "HEADROOM_KOMPRESS_ONNX_INTRA_THREADS" KOMPRESS_ONNX_INTER_THREADS_ENV = "HEADROOM_KOMPRESS_ONNX_INTER_THREADS" KOMPRESS_COREML_CACHE_DIR_ENV = "HEADROOM_KOMPRESS_COREML_CACHE_DIR" KOMPRESS_MAX_CONCURRENT_ENV = "HEADROOM_KOMPRESS_MAX_CONCURRENT" KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV = "HEADROOM_KOMPRESS_EXECUTION_TIMEOUT_MS" KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_DEFAULT = 25 KOMPRESS_BATCH_SIZE_ENV = "HEADROOM_KOMPRESS_BATCH_SIZE" KompressBackend = Literal["auto", "onnx", "onnx_cpu", "onnx_coreml", "pytorch", "pytorch_mps"] # HuggingFace local-lookup errors that mean "asset not in cache" rather than a # genuine failure. Caught when loading cache-only so startup can defer instead. try: from huggingface_hub.errors import EntryNotFoundError, LocalEntryNotFoundError _NOT_CACHED_ERRORS: tuple[type[BaseException], ...] = ( LocalEntryNotFoundError, EntryNotFoundError, OSError, ) except Exception: # pragma: no cover - huggingface_hub always present with [ml] _NOT_CACHED_ERRORS = (OSError,) class KompressModelNotCached(RuntimeError): """Raised when a cache-only load is requested but the model is not cached. Used by startup eager-preload (``allow_download=False``) so the caller can defer the download to first use instead of blocking the proxy startup path on a network fetch. """ # Model cache: model_id -> (model, tokenizer, backend) # Supports multiple models loaded simultaneously. _kompress_cache: dict[str, tuple[Any, Any, str]] = {} _kompress_lock = threading.Lock() _execution_semaphores: dict[str, threading.BoundedSemaphore] = {} _execution_semaphores_lock = threading.Lock() _execution_metrics_lock = threading.Lock() _execution_skip_counters: dict[str, int] = { "timeout": 0, } _execution_wait_seconds_total: dict[str, float] = { "timeout": 0.0, } def _execution_wait_budget_seconds() -> float: raw = os.environ.get(KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV) if raw is None: return KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_DEFAULT / 1000.0 try: parsed = int(raw) except ValueError: logger.warning( "Invalid %s=%r; using %dms", KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV, raw, KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_DEFAULT, ) return KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_DEFAULT / 1000.0 if parsed < 0: logger.warning( "Negative %s=%r; disabling timeout and using fail-open.", KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV, raw, ) return 0.0 return parsed / 1000.0 def _acquire_execution_slot( backend: str, device_type: str, *, timeout_seconds: float | None, ) -> tuple[threading.BoundedSemaphore | None, float]: semaphore = _execution_semaphore(backend, device_type) start = time.perf_counter() if timeout_seconds is None: semaphore.acquire() wait_ms = (time.perf_counter() - start) * 1000.0 return semaphore, wait_ms acquired = semaphore.acquire(blocking=False) if not acquired and timeout_seconds > 0: acquired = semaphore.acquire(timeout=timeout_seconds) elif not acquired and timeout_seconds == 0: acquired = False wait_ms = (time.perf_counter() - start) * 1000.0 if not acquired: with _execution_metrics_lock: _execution_skip_counters["timeout"] += 1 _execution_wait_seconds_total["timeout"] += wait_ms / 1000.0 return None, wait_ms return semaphore, wait_ms def get_kompress_execution_stats() -> dict[str, int | float]: """Return execution-acquire observability counters.""" with _execution_metrics_lock: return { "execution_acquire_timeout_ms": int(_execution_wait_budget_seconds() * 1000), "execution_timeout_skips_total": _execution_skip_counters["timeout"], "execution_wait_seconds_total": _execution_wait_seconds_total["timeout"], } def _selected_backend() -> KompressBackend: raw = os.environ.get(KOMPRESS_BACKEND_ENV, "auto").strip().lower().replace("-", "_") aliases = { "": "auto", "cpu": "onnx_cpu", "coreml": "onnx_coreml", "mps": "pytorch_mps", "torch": "pytorch", "torch_mps": "pytorch_mps", "onnx": "onnx", "onnx_cpu": "onnx_cpu", "onnx_coreml": "onnx_coreml", "pytorch": "pytorch", "pytorch_mps": "pytorch_mps", "auto": "auto", } backend = aliases.get(raw) if backend is None: logger.warning( "%s has unrecognized value %r; falling back to 'auto'. Valid values: %s", KOMPRESS_BACKEND_ENV, os.environ.get(KOMPRESS_BACKEND_ENV, ""), ", ".join(sorted(set(aliases.values()))), ) return "auto" return backend # type: ignore[return-value] def _env_int(name: str) -> int | None: raw = os.environ.get(name) if raw is None or raw.strip() == "": return None try: value = int(raw) except ValueError: logger.warning("%s must be an integer, got %r; ignoring", name, raw) return None if value <= 0: logger.warning("%s must be positive, got %r; ignoring", name, raw) return None return value def _onnx_session_options(ort: Any) -> Any: return create_cpu_session_options( ort, intra_op_num_threads=_env_int(KOMPRESS_ONNX_INTRA_THREADS_ENV), inter_op_num_threads=_env_int(KOMPRESS_ONNX_INTER_THREADS_ENV), ) def _model_device_type(model: Any, backend: str) -> str: if backend.startswith("onnx"): return backend if hasattr(model, "parameters"): try: return str(next(model.parameters()).device.type) except Exception: return "unknown" return "unknown" def _default_max_concurrent(backend: str, device_type: str) -> int: # MPS/CUDA execution is usually serialized under the hood; letting many # Codex unit workers call the same model concurrently mostly adds queueing, # memory pressure, and timeout leaks. CPU defaults to 1 as well because ONNX # already owns its intra/inter-op threads. if backend.startswith("onnx"): return 1 if backend == "pytorch" and device_type in {"cuda", "mps", "cpu"}: return 1 return 1 def _execution_limit(backend: str, device_type: str) -> int: return _env_int(KOMPRESS_MAX_CONCURRENT_ENV) or _default_max_concurrent(backend, device_type) def _execution_semaphore(backend: str, device_type: str) -> threading.BoundedSemaphore: limit = _execution_limit(backend, device_type) key = f"{backend}:{device_type}:{limit}" with _execution_semaphores_lock: semaphore = _execution_semaphores.get(key) if semaphore is None: semaphore = threading.BoundedSemaphore(limit) _execution_semaphores[key] = semaphore return semaphore def _batch_size() -> int: return _env_int(KOMPRESS_BATCH_SIZE_ENV) or 32 def _bucket_count(value: int) -> str: """Return a coarse, privacy-preserving size bucket.""" if value <= 0: return "0" lower = 1 << (value.bit_length() - 1) upper = lower << 1 return f"{lower}-{upper}" def _kompress_content_signature(content: str) -> Any: """Create a first-class TOIN signature for Kompress/plain-text content. This intentionally keys on shape, not values. Retrieval pressure should teach TOIN about this class of compressed content without storing the content or treating it as an anonymous fallback. """ from ..telemetry.models import ToolSignature words = content.split() line_count = content.count("\n") + 1 if content else 0 nonempty_lines = [line for line in content.splitlines() if line.strip()] avg_line_chars = ( sum(len(line) for line in nonempty_lines) // len(nonempty_lines) if nonempty_lines else 0 ) has_paths = "/" in content or "\\" in content has_assignment_like_tokens = any("=" in word for word in words[:200]) has_brackets = any(ch in content for ch in "{}[]()") has_error_terms = any( term in content.lower() for term in ("error", "exception", "traceback", "failed", "fatal") ) shape = "|".join( ( "kompress-text", f"chars:{_bucket_count(len(content))}", f"words:{_bucket_count(len(words))}", f"lines:{_bucket_count(line_count)}", f"avg_line:{_bucket_count(avg_line_chars)}", f"paths:{int(has_paths)}", f"assign:{int(has_assignment_like_tokens)}", f"brackets:{int(has_brackets)}", f"errors:{int(has_error_terms)}", ) ) structure_hash = hashlib.sha256(shape.encode()).hexdigest()[:24] return ToolSignature( structure_hash=structure_hash, field_count=0, has_nested_objects=False, has_arrays=False, max_depth=0, string_field_count=1, has_error_like_field=has_error_terms, has_message_like_field=True, ) def _is_onnx_available() -> bool: """Check if ONNX Runtime is available (lightweight, no torch needed).""" try: import onnxruntime # noqa: F401 import transformers # noqa: F401 return True except ImportError: return False def _is_pytorch_available() -> bool: """Check if full PyTorch stack is available (requires [ml] extra).""" try: import safetensors # noqa: F401 import torch # noqa: F401 import transformers # noqa: F401 return True except ImportError: return False def is_kompress_available() -> bool: """Check if Kompress can run — ONNX (lightweight) or PyTorch (full).""" return _is_onnx_available() or _is_pytorch_available() # ── Model Architecture (must match training) ────────────────────────── # torch/transformers are imported lazily — only when actually needed. # This allows `from kompress_compressor import is_kompress_available` # to work without torch installed. def _get_model_class() -> type: """Return the HeadroomCompressorModel class, importing torch on demand.""" import torch import torch.nn as nn from transformers import AutoModel 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_keep_mask( self, input_ids: torch.Tensor, attention_mask: torch.Tensor ) -> torch.Tensor: """Get per-token keep/discard decision. True = keep.""" with torch.no_grad(): hidden = self.encoder(input_ids, attention_mask=attention_mask).last_hidden_state # Token head: binary classifier — argmax decides keep/discard token_logits = self.token_head(hidden) # [B, L, 2] token_keep = ( token_logits[:, :, 1] > token_logits[:, :, 0] ) # True if class 1 > class 0 # Span head: boost tokens in important spans # If a token is borderline but its span is important, keep it span_scores = self.span_conv(hidden.transpose(1, 2)).squeeze(1) span_boost = span_scores > 0.5 # span says this region matters # Keep if: token head says keep, OR token is borderline and span says keep token_probs = torch.softmax(token_logits, dim=-1)[:, :, 1] borderline = (token_probs > 0.3) & (token_probs <= 0.5) keep = token_keep | (borderline & span_boost) return keep # type: ignore[no-any-return] def get_scores(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: """Get per-token importance scores (for ranking when target_ratio is set).""" 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) # type: ignore[no-any-return] return HeadroomCompressorModel # ── Model Loading ───────────────────────────────────────────────────── class _OnnxModel: """Thin wrapper so ONNX session has the same interface as PyTorch model.""" def __init__(self, session: Any): self._session = session def get_scores(self, input_ids: Any, attention_mask: Any) -> Any: """Return [batch, seq] scores via ONNX Runtime.""" import numpy as np scores = self._session.run( ["final_scores"], { "input_ids": np.asarray(input_ids, dtype=np.int64), "attention_mask": np.asarray(attention_mask, dtype=np.int64), }, ) return scores[0] # [batch, seq] numpy array def get_keep_mask(self, input_ids: Any, attention_mask: Any) -> Any: """Return [batch, seq] boolean mask (score > 0.5).""" import numpy as np scores = self.get_scores(input_ids, attention_mask) return (np.array(scores) > 0.5).tolist() def _onnx_filename_candidates() -> tuple[str, ...]: """ONNX repo paths to try, honoring an optional exact-file override.""" override = os.environ.get(KOMPRESS_ONNX_FILENAME_ENV, "").strip() if override: # Put the override first but keep the defaults as a safety net. return (override, *(f for f in _DEFAULT_ONNX_FILENAMES if f != override)) return _DEFAULT_ONNX_FILENAMES def _create_onnx_session( model_id: str, providers: list[Any], *, allow_download: bool = True ) -> Any: """Resolve and load the model's ONNX artifact, trying candidates in order. A candidate is skipped on download miss (file not in the repo) or on session-load failure (e.g. the weight-only int8 artifact uses the MatMulNBits contrib op, which old onnxruntime builds can't run — those installs fall through to the fp32 artifact instead of losing Kompress). When ``allow_download`` is ``False`` candidates are resolved from the local cache only; if none is cached, :class:`KompressModelNotCached` is raised instead of hitting the network. ``onnxruntime`` is imported only after a candidate resolves, so a cache-only miss never requires it. """ last_err: Exception | None = None cache_miss = False ort: Any = None for filename in _onnx_filename_candidates(): try: onnx_path = hf_hub_download_local_first( model_id, filename, allow_network=allow_download ) except Exception as exc: last_err = exc cache_miss = cache_miss or isinstance(exc, _NOT_CACHED_ERRORS) logger.debug("ONNX artifact %r unavailable for %s: %s", filename, model_id, exc) continue if ort is None: import onnxruntime ort = onnxruntime try: return ort.InferenceSession( onnx_path, _onnx_session_options(ort), providers=providers, ) except Exception as exc: last_err = exc logger.warning( "ONNX artifact %r from %s failed to load (%s); trying next candidate", filename, model_id, exc, ) if not allow_download and cache_miss: raise KompressModelNotCached(model_id) from last_err raise FileNotFoundError( f"No loadable ONNX artifact in {model_id}; tried {_onnx_filename_candidates()}" ) from last_err def _load_kompress_onnx( model_id: str, *, use_coreml: bool = False, allow_download: bool = True, ) -> tuple[Any, Any, str]: """Download ONNX INT8 model from HuggingFace and load with onnxruntime. When ``allow_download`` is ``False`` the model and tokenizer are loaded from the local cache only; a cache miss raises :class:`KompressModelNotCached` instead of hitting the network. """ with _kompress_lock: if model_id in _kompress_cache: return _kompress_cache[model_id] logger.info("Downloading Kompress ONNX model from %s ...", model_id) backend = "onnx_coreml" if use_coreml else "onnx" providers: list[Any] if use_coreml: from headroom import paths as _paths coreml_cache_dir = os.environ.get(KOMPRESS_COREML_CACHE_DIR_ENV, "").strip() cache_dir = ( coreml_cache_dir if coreml_cache_dir else str(_paths.workspace_dir() / "cache" / "coreml") ) os.makedirs(cache_dir, exist_ok=True) providers = [ ( "CoreMLExecutionProvider", { "ModelFormat": "NeuralNetwork", "MLComputeUnits": "ALL", "RequireStaticInputShapes": "1", "ModelCacheDirectory": cache_dir, }, ), "CPUExecutionProvider", ] else: providers = ["CPUExecutionProvider"] session = _create_onnx_session(model_id, providers, allow_download=allow_download) model = _OnnxModel(session) from transformers import AutoTokenizer tokenizer = _load_modernbert_tokenizer(AutoTokenizer, allow_download=allow_download) _kompress_cache[model_id] = (model, tokenizer, backend) logger.info("Kompress ONNX loaded: %s backend=%s", model_id, backend) return model, tokenizer, backend def _load_modernbert_tokenizer(auto_tokenizer: Any, *, allow_download: bool) -> Any: """Load the ModernBERT tokenizer, cache-only when ``allow_download`` is False.""" try: return auto_tokenizer.from_pretrained( "answerdotai/ModernBERT-base", local_files_only=not allow_download ) except _NOT_CACHED_ERRORS as exc: if not allow_download: raise KompressModelNotCached("answerdotai/ModernBERT-base") from exc raise def _load_kompress_pytorch( model_id: str, device: str = "auto", *, allow_download: bool = True ) -> tuple[Any, Any, str]: """Download PyTorch model from HuggingFace and load with torch. When ``allow_download`` is ``False`` weights and tokenizer are loaded from the local cache only; a cache miss raises :class:`KompressModelNotCached`. """ import torch from transformers import AutoTokenizer with _kompress_lock: if model_id in _kompress_cache: return _kompress_cache[model_id] logger.info("Downloading Kompress PyTorch model from %s ...", model_id) try: weights_path = hf_hub_download_local_first( model_id, "model.safetensors", allow_network=allow_download ) except _NOT_CACHED_ERRORS as exc: if not allow_download: raise KompressModelNotCached(model_id) from exc raise HeadroomCompressorModel = _get_model_class() model = HeadroomCompressorModel() from safetensors.torch import load_file state_dict = load_file(weights_path) model.load_state_dict(state_dict, strict=False) 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() tokenizer = _load_modernbert_tokenizer(AutoTokenizer, allow_download=allow_download) _validate_pytorch_device(model, tokenizer, device) _kompress_cache[model_id] = (model, tokenizer, "pytorch") logger.info("Kompress PyTorch loaded on %s (%s)", device, model_id) return model, tokenizer, "pytorch" def _validate_pytorch_device(model: Any, tokenizer: Any, device: str) -> None: if device == "cpu": return encoding = tokenizer( ["headroom", "kompress", "probe"], is_split_into_words=True, truncation=True, max_length=512, padding=True, return_tensors="pt", ) input_ids = encoding["input_ids"].to(device) attention_mask = encoding["attention_mask"].to(device) semaphore, _wait_ms = _acquire_execution_slot( "pytorch", device, timeout_seconds=None, ) assert semaphore is not None with contextlib.ExitStack() as stack: stack.callback(semaphore.release) scores = model.get_scores(input_ids, attention_mask) _ = scores[0].detach().cpu() def _load_kompress( model_id: str = HF_MODEL_ID, device: str = "auto", *, allow_download: bool = True ) -> tuple[Any, Any, str]: """Load Kompress model, returns (model, tokenizer, backend). The default keeps the historic behavior: try ONNX CPU first (lightweight), then fall back to PyTorch. Operators can override via HEADROOM_KOMPRESS_BACKEND: - auto: ONNX CPU first, then PyTorch. - onnx / onnx_cpu: force ONNX CPU. - onnx_coreml: force ONNX Runtime CoreML provider with CPU fallback. - pytorch: force PyTorch with the configured device. - pytorch_mps: force PyTorch on Apple's MPS backend. When ``allow_download`` is ``False`` the model is loaded from the local cache only and a cache miss raises :class:`KompressModelNotCached` rather than fetching from the network. Models are cached by model_id — multiple models can coexist. """ if model_id in _kompress_cache: return _kompress_cache[model_id] backend = _selected_backend() if backend in ("onnx", "onnx_cpu"): return _load_kompress_onnx(model_id, use_coreml=False, allow_download=allow_download) if backend == "onnx_coreml": return _load_kompress_onnx(model_id, use_coreml=True, allow_download=allow_download) if backend in ("pytorch", "pytorch_mps"): forced_device = "mps" if backend == "pytorch_mps" else device try: return _load_kompress_pytorch(model_id, forced_device, allow_download=allow_download) except KompressModelNotCached: raise except Exception as exc: if backend != "pytorch_mps": raise logger.warning( "Kompress PyTorch MPS validation failed for %s; falling back to ONNX CPU: %s", model_id, exc, ) if _is_onnx_available(): return _load_kompress_onnx( model_id, use_coreml=False, allow_download=allow_download ) return _load_kompress_pytorch(model_id, "cpu", allow_download=allow_download) # Auto mode: preserve stable default behavior. This avoids changing # compression quality/perf characteristics for existing installs while # allowing opt-in MPS/CoreML experiments via HEADROOM_KOMPRESS_BACKEND. if _is_onnx_available(): try: return _load_kompress_onnx(model_id, use_coreml=False, allow_download=allow_download) except KompressModelNotCached: # Cache-only miss: don't trigger a PyTorch network download as a # fallback — propagate so the caller can defer. if not allow_download: raise except Exception as e: logger.warning("ONNX load failed for %s, trying PyTorch: %s", model_id, e) if _is_pytorch_available(): return _load_kompress_pytorch(model_id, device, allow_download=allow_download) raise ImportError( "Kompress requires onnxruntime or torch. Install with: pip install headroom-ai[proxy]" ) def unload_kompress_model(model_id: str | None = None) -> bool: """Unload Kompress model(s) to free memory. Args: model_id: Specific model to unload. If None, unloads all cached models. """ with _kompress_lock: if model_id is not None: if model_id in _kompress_cache: del _kompress_cache[model_id] else: return False elif _kompress_cache: _kompress_cache.clear() else: return False try: import torch if torch.cuda.is_available(): torch.cuda.empty_cache() except ImportError: pass gc.collect() trim_process_heap() return True # ── Background model download ───────────────────────────────────────── # # The proxy request path must never block on a cold model download. A first # deep-path request would otherwise resolve the 274MB ONNX artifact via an # inline hf_hub_download on the request thread, where it races the proxy's # compression timeout (HEADROOM_COMPRESSION_TIMEOUT_SECONDS, default 30s). The # fetch is cancelled mid-transfer, the blob never finalizes in the HF cache, # and every subsequent request re-hangs and fails open. Instead the request # path resolves the model cache-only (allow_download=False) and pulls it down # once here, in a daemon thread that the compression timeout does not bound. _download_threads: dict[str, threading.Thread] = {} _download_threads_lock = threading.Lock() def _background_download(model_id: str, device: str) -> None: try: logger.info("Kompress: downloading model %s in the background ...", model_id) _load_kompress(model_id, device, allow_download=True) logger.info("Kompress: background model download complete for %s", model_id) except Exception as exc: logger.warning("Kompress: background model download failed for %s: %s", model_id, exc) def ensure_background_download(model_id: str = HF_MODEL_ID, device: str = "auto") -> None: """Start a one-shot background download of the model if it isn't cached. Idempotent and non-blocking: at most one download thread runs per model_id, and a finished or failed thread is replaced on the next call so a transient network failure can be retried by a later request. Once the download completes the deep path activates on subsequent requests without ever blocking one on the network. """ if model_id in _kompress_cache: return with _download_threads_lock: if model_id in _kompress_cache: return existing = _download_threads.get(model_id) if existing is not None and existing.is_alive(): return thread = threading.Thread( target=_background_download, args=(model_id, device), name=f"kompress-download-{model_id.replace('/', '-')}", daemon=True, ) _download_threads[model_id] = thread thread.start() # ── Compressor ──────────────────────────────────────────────────────── @dataclass class KompressConfig: """Configuration for Kompress compression. The model_id, chunk_words, and score_threshold are coupled: a model trained on 50-word chunks needs chunk_words=50 at inference. The defaults match kompress-v2-base. For domain-specific models, set all three. Example — financial documents:: KompressConfig( model_id="chopratejas/kompress-finance", chunk_words=50, score_threshold=0.5, ) """ device: str = "auto" enable_ccr: bool = True model_id: str = HF_MODEL_ID chunk_words: int = 350 score_threshold: float = 0.5 @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. Auto-downloads the model from HuggingFace on first use. Configure via KompressConfig to select model, chunk size, and threshold. """ name: str = "kompress_compressor" def __init__(self, config: KompressConfig | None = None): self.config = config or KompressConfig() def preload(self, *, allow_download: bool = True) -> str: """Load the backing model/tokenizer and return the selected backend. When ``allow_download`` is ``False`` the model is loaded from the local cache only; if it is not cached, :class:`KompressModelNotCached` is raised so the caller can defer the download to first use. Startup eager preload uses this so a cold cache cannot block the proxy from binding its port. """ _model, _tokenizer, backend = _load_kompress( self.config.model_id, self.config.device, allow_download=allow_download ) return backend def is_ready(self) -> bool: """True if the model is loaded so :meth:`compress` won't touch the network. A plain cache-membership check — no lock, no I/O — safe to call on the hot request path to decide whether to run the deep compressor or skip it. """ return self.config.model_id in _kompress_cache def ensure_background_load(self) -> None: """Kick off a one-shot, non-blocking background download of the model. No-op when the model is already cached or a download is already running. """ ensure_background_download(self.config.model_id, self.config.device) def compress( self, content: str, context: str = "", content_type: str | None = None, question: str | None = None, target_ratio: float | None = None, *, allow_download: bool = True, ) -> KompressResult: """Compress content using Kompress model. Args: content: Text to compress. context: Optional surrounding context (unused by model). content_type: Ignored — model decides importance per content type. question: Ignored — reserved for future QA-aware compression. 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. allow_download: When False, load the model from the local cache only; a cache miss passes through instead of fetching from the network. The proxy sets this False so a cold model never blocks the request thread (see ``ensure_background_download``); direct callers keep the historic auto-download-on-first-use behavior. Returns: KompressResult with compressed text. """ words = content.split() n_words = len(words) if n_words < 10: return self._passthrough(content, n_words) # Cooperative wall-clock budget (#1171): kompress ONNX inference is # O(tokens) and non-preemptible once the request's asyncio timeout fires, # so one large block can run for minutes holding a worker (the leak -> # executor-saturation -> queue-timeout cascade). Bail at the next chunk # boundary past this budget, keeping the unprocessed tail verbatim. 0 # disables. Env HEADROOM_COMPRESSION_DEADLINE_MS overrides (default 20s). # Cached per instance: operator config, read once -- not per compress() call. deadline_s = getattr(self, "_deadline_s", None) if deadline_s is None: try: deadline_s = max( 0.0, float(os.environ.get("HEADROOM_COMPRESSION_DEADLINE_MS", "20000")) / 1000.0, ) except ValueError: deadline_s = 20.0 self._deadline_s = deadline_s try: model, tokenizer, backend = _load_kompress( self.config.model_id, self.config.device, allow_download=allow_download ) is_onnx = backend == "onnx" device_type = _model_device_type(model, backend) if self._should_batch_single_content(model, backend): batch_result = self.compress_batch( [content], context=context, content_type=content_type, question=question, target_ratio=[target_ratio], batch_size=_batch_size(), ) if batch_result: return batch_result[0] max_chunk_words = self.config.chunk_words kept_ids: set[int] = set() inference_ms = 0.0 chunk_count = 0 t_deadline = time.perf_counter() for chunk_start in range(0, n_words, max_chunk_words): if deadline_s and (time.perf_counter() - t_deadline) > deadline_s: # Keep everything from here on verbatim and stop: a partial # compression that returns NOW beats a full one that leaks a # non-preemptible worker for minutes (#1171). kept_ids.update(range(chunk_start, n_words)) logger.warning( "Kompress hit %.1fs deadline after %d/%d words (%d chunks done); " "kept remainder verbatim to free the request thread (#1171)", deadline_s, chunk_start, n_words, chunk_count, ) break chunk_count += 1 chunk_words = words[chunk_start : chunk_start + max_chunk_words] # ONNX uses numpy tensors, PyTorch uses torch tensors return_tensors = "np" if is_onnx else "pt" encoding = tokenizer( chunk_words, is_split_into_words=True, truncation=True, max_length=512, padding=True, return_tensors=return_tensors, ) input_ids = encoding["input_ids"] attention_mask = encoding["attention_mask"] word_ids = encoding.word_ids(batch_index=0) if not is_onnx: device = next(model.parameters()).device input_ids = input_ids.to(device) attention_mask = attention_mask.to(device) semaphore, _wait_ms = _acquire_execution_slot( backend, device_type, timeout_seconds=_execution_wait_budget_seconds(), ) if semaphore is None: logger.warning( "Kompress execution saturated after %.2fms; skipping chunk=%d " "for backend=%s device=%s after deadline path", _wait_ms, chunk_start, backend, device_type, ) return self._passthrough(content, n_words) with contextlib.ExitStack() as stack: stack.callback(semaphore.release) inference_started = time.perf_counter() if target_ratio is not None: scores = model.get_scores(input_ids, attention_mask) if is_onnx: score_list = scores[0] # numpy: [seq_len] else: score_list = scores[0].cpu() else: keep_mask = model.get_keep_mask(input_ids, attention_mask) if is_onnx: mask_list = keep_mask[0] # list of bools else: mask_list = keep_mask[0].cpu() inference_ms += (time.perf_counter() - inference_started) * 1000 if target_ratio is not None: word_scores: dict[int, float] = {} for idx, wid in enumerate(word_ids): if wid is None: continue s = float(score_list[idx]) if wid not in word_scores or s > word_scores[wid]: word_scores[wid] = s if word_scores: sorted_wids = sorted( word_scores, key=lambda w: word_scores[w], reverse=True ) num_keep = max(1, int(len(sorted_wids) * target_ratio)) for wid in sorted_wids[:num_keep]: kept_ids.add(wid + chunk_start) else: for idx, wid in enumerate(word_ids): if wid is None: continue if bool(mask_list[idx]): kept_ids.add(wid + chunk_start) # Hard override: always keep must-keep tokens regardless of model score. # Numbers, error names, paths, and flags carry meaning agents cannot # reconstruct from context. Disable via HEADROOM_KOMPRESS_MUST_KEEP=0. _add_kompress_must_keep_words(kept_ids, chunk_words, chunk_start) if not kept_ids: if inference_ms >= 1000.0: logger.info( "Kompress slow passthrough backend=%s device=%s words=%d chunks=%d " "inference_ms=%.0f", backend, device_type, n_words, chunk_count, inference_ms, ) return self._passthrough(content, n_words) 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, model_used=self.config.model_id, ) # 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}]" ) if inference_ms >= 1000.0: logger.info( "Kompress slow compress backend=%s device=%s words=%d chunks=%d " "inference_ms=%.0f ratio=%.3f saved=%d", backend, device_type, n_words, chunk_count, inference_ms, ratio, result.tokens_saved, ) return result except KompressModelNotCached: logger.debug( "Kompress model %s not cached; passing through without compression", self.config.model_id, ) return self._passthrough(content, n_words) except Exception as e: logger.warning("Kompress compression failed: %s", e) return self._passthrough(content, n_words) def compress_batch( self, contents: list[str], context: str = "", content_type: str | None = None, question: str | None = None, target_ratio: float | list[float | None] | None = None, batch_size: int = 32, ) -> list[KompressResult]: """Compress multiple texts. Uses batched inference on GPU, sequential on CPU. On GPU (PyTorch + CUDA / MPS), runs a single batched forward pass per chunk batch, amortizing model inference across N texts. On CPU (ONNX or PyTorch), falls back to sequential ``compress()`` calls because ONNX Runtime's CPU provider does not parallelize across the batch dimension for this model (empirically 0.7-0.9x vs sequential). The fallback is transparent: callers get the best available performance per device without needing to detect the backend themselves. Measured performance (RTX 3080 Ti, ~350-word inputs): GPU batched vs sequential: N=3: 1.76x speedup N=5: 2.08x speedup N=12: 2.18x speedup N=24: 2.34x speedup CPU (ONNX, 16 logical threads): falls back to sequential; net effect is parity with direct ``compress()`` in a loop. Args: contents: List of texts to compress. May contain short texts or empty strings — those pass through without a model call. context: Unused (parity with ``compress``). content_type: Unused (parity with ``compress``). question: Unused (parity with ``compress``). target_ratio: Compression target, one of: * ``None`` — model decides per text (same as :meth:`compress`). * ``float`` — applied uniformly to every text in the batch. * ``list`` of ``float | None`` — per-text ratio; must match ``len(contents)``. ``None`` entries let the model decide for that text. batch_size: Maximum number of chunks per forward pass on the batched path (GPU only — ignored on CPU fallback). Default ``32`` is a reasonable balance for ModernBERT on GPU. Returns: List of :class:`KompressResult`, one per input text, in input order. Empty input returns empty list. Failed texts fall back to passthrough rather than raising. Notes: On the batched GPU path, scoring uses ``get_scores`` uniformly (threshold at 0.5 when ``target_ratio`` is ``None``). This matches the ONNX non-batched behavior exactly. The PyTorch non-batched path applies an additional borderline + span-boost rule, so results may differ by a small fraction of tokens on ``target_ratio=None`` calls via the batched path vs direct :meth:`compress` on PyTorch. Call :meth:`compress` directly if the exact PyTorch borderline behavior is required. """ n = len(contents) if n == 0: return [] # Normalize target_ratio to a per-text list if isinstance(target_ratio, list): if len(target_ratio) != n: raise ValueError( f"target_ratio list length {len(target_ratio)} does not match " f"contents length {n}" ) ratios: list[float | None] = list(target_ratio) else: ratios = [target_ratio] * n # Fast path: on backends where batch-dim parallelism does NOT help # (ONNX CPU, PyTorch CPU), fall back to sequential `compress()` # internally. This keeps the public API consistent while avoiding the # per-item slowdown measured on ONNX CPU (~0.7-0.9x vs sequential). # GPU users still benefit from the batched forward pass below. if self._should_use_sequential_fallback(): return [ self.compress( content, context=context, content_type=content_type, question=question, target_ratio=r, ) for content, r in zip(contents, ratios, strict=True) ] results: list[KompressResult | None] = [None] * n word_lists: list[list[str]] = [c.split() for c in contents] # Short texts short-circuit to passthrough — no model call needed. max_chunk_words = self.config.chunk_words chunk_queue: list[tuple[int, int, list[str], float | None]] = [] for i, (words, ratio) in enumerate(zip(word_lists, ratios, strict=True)): if len(words) < 10: results[i] = self._passthrough(contents[i], len(words)) continue for chunk_start in range(0, len(words), max_chunk_words): chunk_words = words[chunk_start : chunk_start + max_chunk_words] chunk_queue.append((i, chunk_start, chunk_words, ratio)) if not chunk_queue: # Every input was short — all passthrough, no model needed. return [r for r in results if r is not None] # Load model once for the whole batch. try: model, tokenizer, backend = _load_kompress(self.config.model_id, self.config.device) except Exception as e: logger.warning("Kompress load failed for batch: %s — passthrough all", e) for i in range(n): if results[i] is None: results[i] = self._passthrough(contents[i], len(word_lists[i])) return [r for r in results if r is not None] is_onnx = backend == "onnx" device_type = _model_device_type(model, backend) kept_ids_per_text: dict[int, set[int]] = {i: set() for i in range(n) if results[i] is None} inference_ms = 0.0 for batch_start in range(0, len(chunk_queue), batch_size): batch = chunk_queue[batch_start : batch_start + batch_size] batch_word_lists = [c[2] for c in batch] try: return_tensors = "np" if is_onnx else "pt" encoding = tokenizer( batch_word_lists, is_split_into_words=True, truncation=True, max_length=512, padding=True, return_tensors=return_tensors, ) input_ids = encoding["input_ids"] attention_mask = encoding["attention_mask"] if not is_onnx: device = next(model.parameters()).device input_ids = input_ids.to(device) attention_mask = attention_mask.to(device) # Single forward pass for all chunks in this batch. semaphore, wait_ms = _acquire_execution_slot( backend, device_type, timeout_seconds=_execution_wait_budget_seconds(), ) if semaphore is None: logger.warning( "Kompress execution saturated at batch start after %.2fms; " "passing through remaining batch inputs", wait_ms, ) for text_idx, _, _, _ in batch: if results[text_idx] is None: results[text_idx] = self._passthrough( contents[text_idx], len(word_lists[text_idx]) ) kept_ids_per_text.pop(text_idx, None) continue with contextlib.ExitStack() as stack: stack.callback(semaphore.release) inference_started = time.perf_counter() scores = model.get_scores(input_ids, attention_mask) inference_ms += (time.perf_counter() - inference_started) * 1000 for batch_idx, (text_idx, chunk_start, chunk_words, ratio) in enumerate(batch): word_ids = encoding.word_ids(batch_index=batch_idx) score_list = scores[batch_idx] if is_onnx else scores[batch_idx].cpu() # Token -> word reduction (max score per word). word_scores: dict[int, float] = {} for idx, wid in enumerate(word_ids): if wid is None: continue s = float(score_list[idx]) if wid not in word_scores or s > word_scores[wid]: word_scores[wid] = s if not word_scores: continue if ratio is not None: # Top-k by score. sorted_wids = sorted( word_scores, key=lambda w: word_scores[w], reverse=True ) num_keep = max(1, int(len(sorted_wids) * ratio)) for wid in sorted_wids[:num_keep]: kept_ids_per_text[text_idx].add(wid + chunk_start) else: # Threshold from config (default 0.5, matches ONNX get_keep_mask). for wid, score in word_scores.items(): if score > self.config.score_threshold: kept_ids_per_text[text_idx].add(wid + chunk_start) _add_kompress_must_keep_words( kept_ids_per_text[text_idx], chunk_words, chunk_start ) except Exception as e: logger.warning( "Kompress batch forward pass failed: %s — passthrough affected texts", e ) for text_idx, _, _, _ in batch: if results[text_idx] is None: results[text_idx] = self._passthrough( contents[text_idx], len(word_lists[text_idx]) ) kept_ids_per_text.pop(text_idx, None) # Reconstruct compressed text for each non-passthrough result. for text_idx, kept_ids in kept_ids_per_text.items(): if results[text_idx] is not None: continue content = contents[text_idx] words = word_lists[text_idx] n_words = len(words) if not kept_ids: results[text_idx] = self._passthrough(content, n_words) continue compressed_words = [words[w] for w in sorted(kept_ids) if w < n_words] compressed = " ".join(compressed_words) compressed_count = len(compressed_words) comp_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=comp_ratio, model_used=self.config.model_id, ) if self.config.enable_ccr and comp_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}]" ) results[text_idx] = result # Safety: every slot must be populated. final: list[KompressResult] = [] for i, r in enumerate(results): if r is None: final.append(self._passthrough(contents[i], len(word_lists[i]))) else: final.append(r) if inference_ms >= 1000.0: total_words = sum(len(words) for words in word_lists) total_saved = sum(r.tokens_saved for r in final) logger.info( "Kompress slow batch backend=%s device=%s items=%d chunks=%d " "batch_size=%d words=%d inference_ms=%.0f saved=%d", backend, device_type, n, len(chunk_queue), batch_size, total_words, inference_ms, total_saved, ) return final def _should_batch_single_content(self, model: Any, backend: str) -> bool: if backend != "pytorch": return False device_type = _model_device_type(model, backend) return device_type in {"cuda", "mps"} def _should_use_sequential_fallback(self) -> bool: """Return True if batched inference wouldn't speed up on this backend. Empirically measured: - ONNX CPU: no batch-dim parallelism; batched is 0.7-0.9x vs sequential. - PyTorch CPU: typically similar (conservative fallback). - PyTorch + CUDA: 2.0-2.3x speedup at N>=3 — use batched path. If the model isn't loaded yet, we trigger loading so the backend is known. This is a no-op if the model is already in cache. """ model_id = self.config.model_id if model_id not in _kompress_cache: try: _load_kompress(model_id, self.config.device) except Exception: return True if model_id not in _kompress_cache: return True model, _tokenizer, backend = _kompress_cache[model_id] if backend == "onnx": return True # ONNX CPU provider doesn't parallelize batch dim if backend == "pytorch": try: import torch if hasattr(model, "parameters"): device = next(model.parameters()).device if device.type in ("cuda", "mps"): return False # GPU/MPS benefits from batching _ = torch except ImportError: return True return True # Conservative default: sequential 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 signature = _kompress_content_signature(original) compressed_tokens = len(compressed.split()) store = get_compression_store() cache_key = store.store( original, compressed, original_tokens=original_tokens, compressed_tokens=compressed_tokens, original_item_count=original_tokens, compressed_item_count=compressed_tokens, tool_signature_hash=signature.structure_hash, compression_strategy="kompress", ) with contextlib.suppress(Exception): from ..telemetry import get_toin get_toin().record_compression( tool_signature=signature, original_count=original_tokens, compressed_count=compressed_tokens, original_tokens=original_tokens, compressed_tokens=compressed_tokens, strategy="kompress", ) return cache_key except Exception: return None