model card: metrics re-measured on the leak-free v2.1 test split

e00cc5e verified 13 days ago

8.38 kB

license: apache-2.0
language: en
library_name: transformers
tags:
  - text-compression
  - token-classification
  - modernbert
  - lora
  - kompress
base_model: answerdotai/ModernBERT-base
pipeline_tag: token-classification

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.

Threshold	keep_rate	must_keep_recall	F1	precision	recall
0.30	0.946 (5% drop)	0.997	0.942	0.920	0.964
0.40	0.910 (9% drop)	0.994	0.933	0.929	0.937
0.50 (default)	0.866 (13% drop)	0.989	0.920	0.940	0.901
0.60	0.818 (18% drop)	0.980	0.902	0.949	0.860
0.70	0.752 (25% drop)	0.953	0.870	0.958	0.798

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)

step	epoch	eval_loss	F1	must_keep_recall	keep_rate	precision
2000	0.76	0.347	0.905	0.9849	0.867	0.868
4000 ← selected (best must_keep_recall)	1.52	0.342	0.905	0.9909	0.900	0.852
6000	2.27	0.338	0.907	0.9865	0.881	0.863

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)

artifact	size	notes
`onnx/kompress-int8-wo.onnx`	274 MB	weight-only int8 (MatMulNBits), the default artifact; needs `onnxruntime>=1.24`
`onnx/kompress-fp32.onnx`	601 MB	lossless reference

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.

chopratejas
/

kompress-v2-base