kompress-v2-base

kompress-v2-base

Extractive prompt compressor for LLM proxies. Predicts a keep/drop label per token; the surviving tokens form a compressed version of the input that preserves meaning while reducing token count. It has no decoder: the output is always a subsequence of the input, in the original order.

Based on ModernBERT-base (150M params) with a LoRA adapter (4.4M trainable params, 2.9%) plus a custom dual head (token classifier + 1-D span conv). Trained on 126,617 accepted Pipeline A+B labels (compressor + faithfulness judge) across 17 sources: narrative, dialog, code, agent traces, healthcare, finance, government, scientific, web, summary, and tool-calling.

Distribution. Public. Apache 2.0, free for commercial use.

Use with Headroom (recommended)

Headroom's proxy loads onnx/kompress-int8-wo.onnx from this repo and runs it on CPU with ONNX Runtime; no torch needed. Point it at this repo:

from headroom.transforms.kompress_compressor import KompressCompressor, KompressConfig

compressor = KompressCompressor(KompressConfig(model_id="chopratejas/kompress-v2-base"))
result = compressor.compress(long_tool_output)
print(result.compressed, result.compression_ratio)

Or serve it as an endpoint from this repository's modal_jobs/modal_serve.py with KOMPRESS_MODEL_ID=chopratejas/kompress-v2-base set at deploy time.

Use directly (PyTorch)

import torch
from transformers import AutoModel, AutoTokenizer
from huggingface_hub import hf_hub_download

ckpt = torch.load(hf_hub_download("chopratejas/kompress-v2-base", "merged.pt"), map_location="cpu", weights_only=False)
encoder = AutoModel.from_pretrained("answerdotai/ModernBERT-base", attn_implementation="eager")
encoder.load_state_dict({k: v.float() for k, v in ckpt["encoder_state_dict"].items()})
H = encoder.config.hidden_size
token_head = torch.nn.Linear(H, 2)
span_conv = torch.nn.Sequential(torch.nn.Conv1d(H, 256, 5, padding=2), torch.nn.GELU(),
                                torch.nn.Conv1d(256, 1, 3, padding=1))
token_head.load_state_dict(ckpt["token_head_state_dict"])
span_conv.load_state_dict(ckpt["span_conv_state_dict"])

tok = AutoTokenizer.from_pretrained("chopratejas/kompress-v2-base")
enc = tok("The quick brown fox jumps over the lazy dog.", return_tensors="pt")
with torch.no_grad():
    h = encoder(**enc).last_hidden_state
    p_keep = torch.softmax(token_head(h), -1)[..., 1]
    span = torch.sigmoid(span_conv(h.transpose(1, 2)).squeeze(1))
    scores = p_keep * (0.5 + 0.5 * span)          # final_scores in [0, 1]
kept = enc["input_ids"][0][scores[0] >= 0.5]
print(tok.decode(kept, skip_special_tokens=True))

model.safetensors is the HF Trainer checkpoint of HeadroomCompressorV2 (PEFT-wrapped encoder + LoRA + heads) and loads through the kompress package; merged.pt is the same model with LoRA folded into the encoder and is what the ONNX exports were traced from.

Threshold tuning

The model emits final_scores ∈ [0, 1] per subword. Raise the threshold to compress harder; lower it to protect must-keep recall.

Measured on the held-out test split (n=12,697 rows, 16,827,490 scored subwords), checkpoint best, PyTorch bf16 on H100.

Test split. dataset_v2_v2.1-f00 test: split by source document, exact and near duplicates of training rows removed. An earlier version of this card reported metrics on the v2.0 split, where 12.2% of test rows were byte-identical to training rows; those numbers were optimistic by roughly 0.02 F1 and have been replaced. At the default threshold (0.5): F1 0.920, must-keep recall 0.989, keep rate 0.866 (13% of subwords removed).

On the subset the labeller actually compressed (it dropped at least 5% of the words, 6,559,681 subwords), at 0.5: F1 0.893, must-keep recall 0.986, keep rate 0.853. The rest of the split is text where almost every word should be kept, which inflates F1 for any model.

Training data

• 126,617 labeled examples after min_drop_ratio=0.05 filtering and same-conversation packing (from 367,525 accepted labels).
• Sources: arxiv, pubmed-scientific, govreport, swe-smith, swe-gym-openhands, toolmind, xlam-fc, fineweb-edu, cnn-dailymail, xsum, glaive-fc, lmsys-chat, claude-code-sessions, meetingbank, the-stack-smol-md, samsum, swe-bench-verified.
• Labeler: DeepSeek-V4-Flash (compressor) + DeepSeek-V4-Pro (judge) with Pipeline A + B faithfulness loop. Hard-keep overlay enforces names, dates, numbers, URLs, code identifiers via GLiNER + regex + lexicons.
• Bucket split: short=48%, mid=31%, long=21% (max_length 8,192 native ModernBERT context).
• Split: train=126,617 / val=7,037 / test=7,037.
• Data card, per-source licences and the two non-commercial sources (4.78% of the corpus): DATA.md and LICENSES.md in the training repository.

Training details

• Base: ModernBERT-base (150M params)
• Encoder fine-tuning: LoRA (r=16, alpha=32, target_modules=Wqkv/Wi/Wo)
• Heads: per-token CE (must-keep loss weight = 3.0) + 1-D span conv (BCE, weight 0.3 on total loss)
• Trainable params: 4.4M (2.9% of total)
• Optim: AdamW (lr=2e-4 cosine, warmup_ratio=0.06, weight_decay=0.01)
• Effective batch: 48; epochs: 3; bf16 with FlashAttention-2 + gradient checkpointing
• Hardware: 1×H100 80GB, ~39 min wall-clock

Validation curve (n=7,037, threshold=0.5)

The shipped weights are the checkpoint with the best validation must-keep recall (metric_for_best_model), evaluated every 2,000 steps.

ONNX (what Headroom loads)

Weight-only int8 agrees with fp32 on 99.6% of keep decisions. Headroom tries int8-wo first and falls back to fp32 on runtimes without the 8-bit kernel.

Files in this repo

config.json                 # KompressV2Config + arch metadata
model.safetensors           # HF Trainer checkpoint (PEFT-wrapped encoder + LoRA + heads)
merged.pt                   # LoRA merged into the encoder; source of the ONNX exports
onnx/kompress-int8-wo.onnx  # weight-only int8, Headroom's default artifact
onnx/kompress-fp32.onnx     # lossless reference
adapter/                    # LoRA adapter + heads only, for stacking per-org adapters
  adapter/adapter_config.json
  adapter/adapter_model.safetensors
  token_head.pt
  span_conv.pt
tokenizer.json              # answerdotai/ModernBERT-base tokenizer
tokenizer_config.json
special_tokens_map.json
README.md                   # this file

License

Apache 2.0 for the model artifact. ModernBERT-base is also Apache 2.0. Training-data licences are per source; see LICENSES.md in the training repository.

See also

• chopratejas/kompress-v2-base — ModernBERT-base variant (149M params, public)
• chopratejas/kompress-v2-large — ModernBERT-large variant (395M params, private/enterprise)
• Headroom proxy integration guide: docs/CUSTOMER_QUICKSTART.md
• Per-org fine-tuning (LoRA stacking): docs/DEPLOYMENT_STRATEGY.md