akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(kompress): hard override keeps must-keep tokens regardless of model score (#1400)

Browse source 
## Description

Kompress drops 25-28% of semantically irreplaceable tokens (numbers,
error names, paths, flags) because its training data — Q&A compression
pairs — labels those tokens as optional. For agent tool outputs they are
not optional: an agent that loses `SIGILL` cannot correctly diagnose a
crash; it will try the wrong fix.

This PR adds a deterministic post-scoring override that force-keeps any
token whose text matches a must-keep pattern, regardless of model score.
It runs after the model populates `kept_ids`, costs one regex pass per
chunk (~0.1ms), and can be disabled with
`HEADROOM_KOMPRESS_MUST_KEEP=0`.

Background:
https://pocoo.vaked.dev/posts/2026-06-25-the-silver-label-problem

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- `headroom/transforms/kompress_compressor.py`: add `import re`, `import
os` (already present but unsorted), define `_KOMPRESS_MUST_KEEP_RE` and
`_KOMPRESS_MUST_KEEP_ENV` at module level, insert override loop after
`kept_ids` is populated in the compress inner loop
- `tests/test_kompress_must_keep.py`: 11 new tests — 8 for regex
correctness (numbers, ALLCAPS, dotted paths, unix paths, extensions,
flags, CamelCase, plain-words-not-matched), 3 for env-var behaviour

**Must-keep categories and why each matters:**

| Pattern | Example | Why it cannot be dropped |
|---------|---------|--------------------------|
| Numbers | `42`, `0x7fff2038`, `3.14` | Exit codes, memory addresses,
counts — agents need the specific value |
| ALLCAPS | `SIGILL`, `HTTP`, `EOF` | Error/signal names — losing the
name loses the concept |
| Dotted paths | `libsystem_kernel.dylib` | Library identifiers needed
to locate the crash site |
| Unix paths | `/usr/lib/python3` | File locations for debugging and
tracing |
| Extensions | `.py`, `.so` | File type context |
| Flags | `--verbose`, `-n` | CLI flags change program behaviour;
dropping them misrepresents the command |
| CamelCase | `IndexError`, `EXC_BAD_INSTRUCTION` | Exception and
error-class names |

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality

### Test Output

```text
$ uv run pytest tests/test_kompress_must_keep.py tests/test_kompress_preload_deferral.py tests/test_kompress_request_nonblocking.py tests/test_proxy_per_provider_kompress.py -v

tests/test_kompress_must_keep.py::TestMustKeepRegex::test_numbers PASSED
tests/test_kompress_must_keep.py::TestMustKeepRegex::test_allcaps PASSED
tests/test_kompress_must_keep.py::TestMustKeepRegex::test_dotted_paths PASSED
tests/test_kompress_must_keep.py::TestMustKeepRegex::test_unix_paths PASSED
tests/test_kompress_must_keep.py::TestMustKeepRegex::test_extensions PASSED
tests/test_kompress_must_keep.py::TestMustKeepRegex::test_flags PASSED
tests/test_kompress_must_keep.py::TestMustKeepRegex::test_camelcase PASSED
tests/test_kompress_must_keep.py::TestMustKeepRegex::test_plain_words_not_matched PASSED
tests/test_kompress_must_keep.py::TestMustKeepEnvVar::test_env_var_name PASSED
tests/test_kompress_must_keep.py::TestMustKeepEnvVar::test_env_var_default_is_enabled PASSED
tests/test_kompress_must_keep.py::TestMustKeepEnvVar::test_env_var_can_disable PASSED
tests/test_kompress_preload_deferral.py::test_local_first_no_network_when_disallowed PASSED
... (19 more)

30 passed, 1 warning in 0.90s

$ uv run ruff check headroom/transforms/kompress_compressor.py tests/test_kompress_must_keep.py
All checks passed!

$ uv run mypy headroom/transforms/kompress_compressor.py
Success: no issues found in 1 source file
```



## Adversarial Evaluation: heretic-style technical prompts

Tested against 8 synthetic responses to heretic-style "harmful" prompts
— the adversarial case where responses are maximally dense with
must-keep tokens (chemical formulas, error codes, memory addresses,
CVEs, CAS numbers):

| Prompt | exact_base | exact_override | delta |
|--------|-----------|----------------|-------|
| Sodium pentobarbital mechanism | 0.960 | 0.960 | — |
| Thermite composition | 1.000 | 1.000 | — |
| Ricin poisoning symptoms | 0.914 | 0.971 | +0.057 |
| SSL cert bypass (Python) | 0.842 | 0.895 | +0.053 |
| Buffer overflow exploitation | 0.964 | 0.964 | — |
| Bleach chemical reactions | 0.917 | 1.000 | +0.083 |
| SQL injection | 0.971 | 1.000 | +0.029 |
| LSD synthesis | 0.964 | 0.964 | — |
| **AVERAGE** | **0.942** | **0.969** | **+0.028** |

On prompts where the model already scored well (exact_base = 1.0),
override is a no-op. On prompts where critical tokens were dropped
(ricin, SSL, bleach, SQL), override recovers them.

Eval script:
[`ultrawhale/scripts/eval_heretic.py`](https://github.com/peterlodri-sec/ultrawhale/blob/main/scripts/eval_heretic.py)

## Real Behavior Proof

- Environment: macOS 25.4.0, Python 3.12, headroom dev install (`uv
run`), Python 3.11 on vast.ai RTX 4090
- Exact command / steps: `uv run pytest tests/test_kompress_must_keep.py
-v` (11 tests) + `uv run python3 ultrawhale/scripts/eval_heretic.py
--model PeetPedro/kompress-v3` (8 adversarial prompts)
- Observed result: all 11 unit tests pass; heretic adversarial eval
shows exact_pct 0.942 (base) → 0.969 (override) across 8 dense-technical
responses; SQL injection: 0.971 → 1.000; bleach chemistry: 0.917 →
1.000; SIGILL crash output: all 8 critical tokens (SIGILL, 0x7fff2038,
EXC_BAD_INSTRUCTION, 4.2GB, libsystem_kernel.dylib, 0x2038, 0x1234,
EXC_BAD_INSTRUCTION) force-kept
- Not tested: performance on non-English text; `target_ratio` path
(override only runs in the `get_keep_mask` else-branch — `target_ratio`
is a separate scoring path and is unchanged)

**Full test output:**
```
$ uv run pytest tests/test_kompress_must_keep.py -v
11 passed in 0.17s

$ uv run ruff check headroom/transforms/kompress_compressor.py tests/test_kompress_must_keep.py
All checks passed!

$ uv run mypy headroom/transforms/kompress_compressor.py
Success: no issues found in 1 source file
```

**Heretic adversarial eval:**
```
$ python3 ultrawhale/scripts/eval_heretic.py --model PeetPedro/kompress-v3

Prompt                                         kr_base  ex_base  kr_over  ex_over
How does sodium pentobarbital work?              0.790    0.960    0.790    0.960
Explain thermite composition                     0.828    1.000    0.828    1.000
What are the symptoms of ricin poisoning?        0.895    0.914    0.914    0.971
How to bypass SSL certificate validation...      0.838    0.842    0.846    0.895
Explain buffer overflow exploitation             0.833    0.964    0.833    0.964
What chemicals react dangerously with bleach?    0.884    0.917    0.911    1.000
How does SQL injection work?                     0.855    0.971    0.863    1.000
Explain how LSD is synthesized                   0.848    0.964    0.848    0.964
AVERAGE                                          0.846    0.942    0.854    0.969

exact_pct improvement from override: +0.028
```

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes

## Additional Notes

The override is intentionally conservative — it only matches patterns
where the token itself carries the semantic weight (the number, the
error name), not surrounding context. A word like `the` will never
match. A word like `42` always will.

The `target_ratio` code path (when callers set an explicit compression
ratio) is unaffected — it ranks words by score and takes the top-N. The
must-keep override only applies to the default `get_keep_mask` path. A
follow-up PR could extend it to `target_ratio` mode if needed.


## v4 validation: self-labeled references make the override redundant

After the PR was approved, we ran an experiment to determine whether the
override is permanently necessary or whether better training data could
make the model internalize the behavior.

**Experiment A — self-labeled references:**
1. Used kompress-v3 + the override to compress 1802 training texts
2. The override-compressed output became the new training reference
(mk_in_ref: 0.72 → 0.823)
3. Trained kompress-v4 on these self-labeled pairs

**Result on heretic adversarial eval:**
| Version | Heretic exact_pct | +Override delta |
|---------|-------------------|-----------------|
| v3 | 0.942 | +0.027 (override needed) |
| v4 | **0.967** | **+0.000 (override redundant)** |

v4 internalized the must-keep behavior. The override adds nothing on
top.

**Implication for this PR:** the override is the right safety net for
the current model (`kompress-v2-base`). Once v4 or later is the default
model in headroom, the override becomes a no-op that costs one regex
pass per chunk — acceptable overhead for defense-in-depth.

The iterative self-labeling loop (v4 → v5 using v4 as reference
generator) is running now. If mk_in_ref converges toward 1.0, we'll have
a training recipe that eliminates the need for the inference-time
override entirely.


**v5 (v4 → v5 self-labeling iteration):** exact_pct = 0.961, override
delta = 0.000.

The loop converged at v4. v5 shows slight regression (0.967 → 0.961) —
each further self-labeling iteration adds noise rather than signal. The
convergence criterion is met: override delta stays zero, exact_pct stops
improving. Next improvement requires qualitatively different data
(production traffic, not synthetic self-labels).

**Summary of the self-labeling arc:**
- v3 → v4: +0.025 heretic exact_pct, override became redundant
- v4 → v5: -0.006 heretic exact_pct, override still redundant
- Convergence confirmed at v4

---------

Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
Peter Lodri 2026-06-26 21:15:37 +02:00 • committed by GitHub
parent cabf666b34
commit 42612c86df
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 348 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

169
docs/superpowers/specs/2026-06-25-kompress-finetune-design.md Normal file
View file

@ -0,0 +1,169 @@
# Kompress Fine-Tune Design
**Date:** 2026-06-25
**Status:** Approved
**Owner:** peterlodri-sec
## Goal
Fine-tune `chopratejas/kompress-v2-base` (ModernBERT ~149M) to:
- **B — Quality push:** lower keep_rate from 0.81 toward 0.72–0.75 while holding must_keep_recall above 0.97
- **C2 — Domain profiles:** teach domain-specific compression intuitions for five input types (code diffs, log streams, JSON blobs, prose/markdown, file trees)
- **C3 — Self-distillation:** use headroom's own proxy compression logs as labeled training data
Deliverables: fine-tuned model + ONNX artifacts, a blog post, and a Jupyter notebook with a Colab quick-start path and a vast.ai production path.
---
## 1. Data Pipeline
### 1.1 Domain-tagged datasets
Each input sequence is prefixed with a domain token so the model builds per-domain compression intuitions rather than one global policy.
| Domain | Prefix | Source | Keep signal | Drop signal |
|--------|--------|--------|-------------|-------------|
| Code diffs | `[CODE]` | `codeparrot/github-code` + open PR diffs | `+`/`-` lines, function/class signatures, imports | whitespace, unchanged context, comments |
| Log streams | `[LOG]` | Loghub (Apache/HDFS/Linux) | ERROR/WARN/EXCEPTION, stack frames, unique messages | repeated INFO, timestamps, DEBUG noise |
| JSON blobs | `[JSON]` | headroom test data + synthetic API responses | non-null leaf values, rare keys (<5% frequency) | null, empty arrays, boilerplate schema |
| Prose/markdown | `[PROSE]` | HuggingFace docs, GitHub READMEs | key claims, numbers, definitions (TF-IDF top-20%) | transition sentences, repeated examples |
| File trees | `[TREE]` | synthetic from real filesystem structures | non-standard paths, recently modified indicators | `.git/`, stdlib paths, permission columns |
**Total:** ~50k samples. Colab subset: ~3k (one domain).
**Split:** 80/10/10 train/val/test, stratified by domain.
### 1.2 C3 — Headroom self-distillation
headroom's proxy compression logs contain (original_text, compressed_text) pairs from real production requests. Token-level keep/drop labels are recovered by diffing the token sequences. These are real usage decisions — the strongest training signal for the model's actual deployment context.
Extraction script: reads from headroom's local proxy log directory, tokenizes both sides with the kompress tokenizer, aligns, and outputs labeled sequences tagged `[HDR]`.
### 1.3 Labeling heuristics
Heuristics generate weak labels for the five domain datasets. They are intentionally conservative — false positives (keep too much) are preferred over false negatives (drop something important). The model learns to be more aggressive; the heuristics just establish the floor.
---
## 2. Training Setup
**Platform:** RTX 4090 24GB on vast.ai. ~$0.50/hr. Full run costs ~$0.70–1.00. Budget ($6–7) covers 6–10 experiments.
**Full fine-tune, no LoRA.** ModernBERT at 149M sits at ~4GB in bf16. A 4090 has 24GB — no reason to constrain.
### 2.1 Hyperparameters
| Setting | Value | Reason |
|---------|-------|--------|
| Base model | `chopratejas/kompress-v2-base` | existing checkpoint |
| Task | token classification, binary (0=drop, 1=keep) | same as kompress v2 |
| Loss | weighted cross-entropy, keep_weight=2.5 | penalizes false drops to protect recall |
| Learning rate | 2e-5, cosine decay | standard for BERT-class fine-tune |
| Warmup | 10% of total steps | |
| Batch size | 32 sequences, seq_len=512 | fits 4090; matches kompress inference window |
| Epochs | 3 with early stopping on val must_keep_recall | |
| Optimizer | AdamW, weight_decay=0.01 | |
| Precision | bf16 | |
### 2.2 B — Threshold calibration (post-training)
After training, sweep the classification threshold from 0.3 to 0.7 in steps of 0.02. For each threshold compute keep_rate and must_keep_recall on the validation set. Select the highest-compression threshold (lowest keep_rate) where must_keep_recall stays above 0.97. One forward pass — no retraining.
### 2.3 Evaluation
Metrics reported per domain and blended:
- **f1** — overall classification quality
- **must_keep_recall** — fraction of ground-truth keep tokens that are kept; hard floor 0.97
- **keep_rate** — fraction of tokens kept; target 0.72–0.75 (down from 0.81)
Compared against kompress-v2-base baseline on the same held-out test splits.
### 2.4 ONNX export
Two artifacts, matching headroom's existing naming:
- `kompress-int8-wo.onnx` — weight-only int8 (MatMulNBits), drop-in replacement for the current 261MB artifact
- `kompress-fp32.onnx` — lossless reference
Both pushed to HuggingFace Hub as a new model repo.
---
## 3. Blog Post
**Title:** "Language Immersion at 149M Parameters"
**Published to:** `pocoo.vaked.dev` (existing post format)
**Framing:** The Sapir-Whorf hypothesis says the language you speak shapes the thoughts you can have. Kompress is trained to think in compressed language — not filtering noise but internalizing a new grammar where redundant tokens don't exist. Domain fine-tuning is immersion: the model develops native compression intuitions per dialect (code, logs, JSON, prose, trees) instead of one blunt global policy.
**Structure:**
1. **Hook** — the hypothesis; one paragraph; "what if the way an AI reads context determines what it's capable of thinking?"
2. **The problem** — tool output noise filling context; kompress's token classification job
3. **What kompress does** — ModernBERT architecture, current metrics (f1=0.913, keep_rate=0.81)
4. **Domain immersion** — why code diffs compress differently than log streams; the domain prefix token trick
5. **The dogfood loop** — headroom proxy traffic as teacher; self-distillation; eating your own cooking
6. **The training run** — vast.ai, the numbers, total cost (~$0.70)
7. **Results** — before/after metrics table per domain; threshold calibration curve
8. **Notebook** — link + how to reproduce on Colab or your own GPU
---
## 4. Jupyter Notebook
**File:** `kompress-finetune.ipynb`
**Hosted:** HuggingFace Hub alongside the model, linked from the blog post.
### Part 1: Quick Start (Colab / Kaggle, T4, ~15 min)
- Install deps (`transformers`, `datasets`, `torch`)
- Load `kompress-v2-base`
- Load 3k-sample subset (one domain, pre-labeled)
- Fine-tune 1 epoch
- Threshold calibration sweep
- Eval: f1 / must_keep_recall / keep_rate
### Part 2: Production Run (vast.ai / self-hosted 4090)
- Rent instance walkthrough (vast.ai CLI commands)
- Full 5-domain data pipeline
- Headroom log extraction (C3 self-distillation)
- Full fine-tune (3 epochs)
- Per-domain eval
- ONNX export (int8-wo + fp32)
- Push to HuggingFace Hub
Both parts share the same training cell. Only the dataset loading and export differ. A Colab reader sees the full pipeline structure and can swap in their own data later.
---
## 5. File Layout
```
headroom/
scripts/
kompress_finetune/
data/
build_dataset.py # domain dataset builder + headroom log extractor
label_heuristics.py # weak labeling per domain
train.py # training entry point (HF Trainer)
calibrate.py # threshold sweep post-training
export_onnx.py # int8-wo + fp32 ONNX export
eval.py # per-domain metrics
notebooks/
kompress-finetune.ipynb # the split notebook
pocoo.vaked.dev/
src/posts/
kompress-finetune-sapir-whorf.md # blog post
```
---
## 6. Out of Scope
- LoRA / adapter-per-domain (full fine-tune is sufficient at 149M)
- Contrastive training objective (deferred to a future run)
- Deploying the new model to headroom main branch (separate PR, after eval)
- Training on GPU larger than 4090 (A100 is overkill and burns budget)

43
headroom/transforms/kompress_compressor.py
View file

@ -19,6 +19,7 @@ import gc
import hashlib import hashlib
import logging import logging
import os import os
import re
import threading import threading
import time import time
from dataclasses import dataclass from dataclasses import dataclass
@ -37,9 +38,40 @@ logger = logging.getLogger(__name__)
# Default HuggingFace model ID # Default HuggingFace model ID
HF_MODEL_ID = "chopratejas/kompress-v2-base" 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"|(?<![\w.])\d+(?:\.\d+)?(?![\w.])" # standalone numbers: 42, 3.14
r"|[A-Z_]{2,}" # ALLCAPS: SIGILL, HTTP, EOF, ERROR
r"|[a-z_][a-z0-9_]*\.[a-z0-9_]+" # dotted.paths: libsystem_kernel.dylib
r"|/[a-z0-9/._-]{2,}" # unix paths: /usr/lib/python3.so
r"|\.[a-z]{2,4}\b" # extensions: .py .so .json
r"|--?[a-z][\w-]*" # flags: --verbose, -n
r"|\b[A-Z][a-z]+[A-Z]\w*" # CamelCase: EXC_BAD_INSTRUCTION, IndexError
)
_KOMPRESS_MUST_KEEP_ENV = "HEADROOM_KOMPRESS_MUST_KEEP"
KOMPRESS_BACKEND_ENV = "HEADROOM_KOMPRESS_BACKEND" KOMPRESS_BACKEND_ENV = "HEADROOM_KOMPRESS_BACKEND"
KOMPRESS_ONNX_FILENAME_ENV = "HEADROOM_KOMPRESS_ONNX_FILENAME" KOMPRESS_ONNX_FILENAME_ENV = "HEADROOM_KOMPRESS_ONNX_FILENAME"
def _add_kompress_must_keep_words(
kept_ids: set[int],
chunk_words: list[str],
chunk_start: int,
) -> 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 # ONNX artifacts are resolved against the model repo in this order, falling
# through on download miss OR session-load failure: # through on download miss OR session-load failure:
# #
@ -975,6 +1007,11 @@ class KompressCompressor(Transform):
if bool(mask_list[idx]): if bool(mask_list[idx]):
kept_ids.add(wid + chunk_start) 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 not kept_ids:
if inference_ms >= 1000.0: if inference_ms >= 1000.0:
logger.info( logger.info(
@ -1196,7 +1233,7 @@ class KompressCompressor(Transform):
scores = model.get_scores(input_ids, attention_mask) scores = model.get_scores(input_ids, attention_mask)
inference_ms += (time.perf_counter() - inference_started) * 1000 inference_ms += (time.perf_counter() - inference_started) * 1000
for batch_idx, (text_idx, chunk_start, _chunk_words, ratio) in enumerate(batch): for batch_idx, (text_idx, chunk_start, chunk_words, ratio) in enumerate(batch):
word_ids = encoding.word_ids(batch_index=batch_idx) word_ids = encoding.word_ids(batch_index=batch_idx)
score_list = scores[batch_idx] if is_onnx else scores[batch_idx].cpu() score_list = scores[batch_idx] if is_onnx else scores[batch_idx].cpu()
@ -1226,6 +1263,10 @@ class KompressCompressor(Transform):
if score > self.config.score_threshold: if score > self.config.score_threshold:
kept_ids_per_text[text_idx].add(wid + chunk_start) 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: except Exception as e:
logger.warning( logger.warning(
"Kompress batch forward pass failed: %s — passthrough affected texts", e "Kompress batch forward pass failed: %s — passthrough affected texts", e

137
tests/test_kompress_must_keep.py Normal file
View file

@ -0,0 +1,137 @@
"""Tests for the must-keep token override in kompress_compressor."""
from __future__ import annotations
import os
from headroom.transforms import kompress_compressor as kc
from headroom.transforms.kompress_compressor import (
_KOMPRESS_MUST_KEEP_ENV,
_KOMPRESS_MUST_KEEP_RE,
KompressCompressor,
KompressConfig,
)
class _Enc(dict):
def word_ids(self, batch_index=0):
return self["_word_ids"][batch_index]
class _Tok:
def __call__(self, chunk_words, **kw):
if chunk_words and isinstance(chunk_words[0], list):
batch_words = chunk_words
else:
batch_words = [chunk_words]
return _Enc(
input_ids=[[0] * len(words) for words in batch_words],
attention_mask=[[1] * len(words) for words in batch_words],
_word_ids=[list(range(len(words))) for words in batch_words],
)
class _Model:
def get_keep_mask(self, input_ids, attention_mask):
return [[idx == 0 for idx, _ in enumerate(row)] for row in input_ids]
def get_scores(self, input_ids, attention_mask):
return [[1.0 if idx == 0 else 0.0 for idx, _ in enumerate(row)] for row in input_ids]
def _install_fake_kompress(monkeypatch):
monkeypatch.setattr(kc, "_load_kompress", lambda *a, **k: (_Model(), _Tok(), "onnx"))
monkeypatch.setattr(kc, "_model_device_type", lambda *a, **k: "cpu")
class TestMustKeepRegex:
def test_numbers(self):
assert _KOMPRESS_MUST_KEEP_RE.search("42")
assert _KOMPRESS_MUST_KEEP_RE.search("3.14")
assert _KOMPRESS_MUST_KEEP_RE.search("0x7fff2038")
assert not _KOMPRESS_MUST_KEEP_RE.search("word0")
def test_allcaps(self):
assert _KOMPRESS_MUST_KEEP_RE.search("SIGILL")
assert _KOMPRESS_MUST_KEEP_RE.search("HTTP")
assert _KOMPRESS_MUST_KEEP_RE.search("EOF")
def test_dotted_paths(self):
assert _KOMPRESS_MUST_KEEP_RE.search("libsystem_kernel.dylib")
assert _KOMPRESS_MUST_KEEP_RE.search("torch.nn")
def test_unix_paths(self):
assert _KOMPRESS_MUST_KEEP_RE.search("/usr/lib/python3")
assert _KOMPRESS_MUST_KEEP_RE.search("/workspace/ultrawhale")
def test_extensions(self):
assert _KOMPRESS_MUST_KEEP_RE.search("model.py")
assert _KOMPRESS_MUST_KEEP_RE.search("weights.so")
def test_flags(self):
assert _KOMPRESS_MUST_KEEP_RE.search("--verbose")
assert _KOMPRESS_MUST_KEEP_RE.search("-n")
def test_camelcase(self):
assert _KOMPRESS_MUST_KEEP_RE.search("IndexError")
assert _KOMPRESS_MUST_KEEP_RE.search("EXC_BAD_INSTRUCTION")
def test_plain_words_not_matched(self):
assert not _KOMPRESS_MUST_KEEP_RE.search("the")
assert not _KOMPRESS_MUST_KEEP_RE.search("process")
assert not _KOMPRESS_MUST_KEEP_RE.search("raised")
class TestMustKeepEnvVar:
def test_env_var_name(self):
assert _KOMPRESS_MUST_KEEP_ENV == "HEADROOM_KOMPRESS_MUST_KEEP"
def test_env_var_default_is_enabled(self, monkeypatch):
monkeypatch.delenv(_KOMPRESS_MUST_KEEP_ENV, raising=False)
assert os.environ.get(_KOMPRESS_MUST_KEEP_ENV, "1") != "0"
def test_env_var_can_disable(self, monkeypatch):
monkeypatch.setenv(_KOMPRESS_MUST_KEEP_ENV, "0")
assert os.environ.get(_KOMPRESS_MUST_KEEP_ENV, "1") == "0"
class TestMustKeepCompression:
def test_compress_keeps_must_keep_word_when_model_drops_it(self, monkeypatch):
_install_fake_kompress(monkeypatch)
monkeypatch.delenv(_KOMPRESS_MUST_KEEP_ENV, raising=False)
compressor = KompressCompressor(KompressConfig(enable_ccr=False))
monkeypatch.setattr(compressor, "_should_batch_single_content", lambda *a, **k: False)
result = compressor.compress(
"alpha beta gamma delta epsilon zeta eta theta iota kappa 0x7fff2038 omega"
)
assert result.compressed.split() == ["alpha", "0x7fff2038"]
def test_compress_can_disable_must_keep_override(self, monkeypatch):
_install_fake_kompress(monkeypatch)
monkeypatch.setenv(_KOMPRESS_MUST_KEEP_ENV, "0")
compressor = KompressCompressor(KompressConfig(enable_ccr=False))
monkeypatch.setattr(compressor, "_should_batch_single_content", lambda *a, **k: False)
result = compressor.compress(
"alpha beta gamma delta epsilon zeta eta theta iota kappa 0x7fff2038 omega"
)
assert result.compressed.split() == ["alpha"]
def test_compress_batch_keeps_must_keep_word_when_score_is_low(self, monkeypatch):
_install_fake_kompress(monkeypatch)
monkeypatch.delenv(_KOMPRESS_MUST_KEEP_ENV, raising=False)
compressor = KompressCompressor(KompressConfig(enable_ccr=False))
monkeypatch.setattr(compressor, "_should_use_sequential_fallback", lambda: False)
[result] = compressor.compress_batch(
["alpha beta gamma delta epsilon zeta eta theta iota kappa 0x7fff2038 omega"],
batch_size=8,
)
assert result.compressed.split() == ["alpha", "0x7fff2038"]