Cut Your LLM Bills in Half: Building a ModernBERT-Based Prompt Compressor (Zero Extra API Calls)

# compressor.py
import math
import torch
from collections import Counter
from transformers import AutoTokenizer, AutoModel

MODEL_ID = "answerdotai/ModernBERT-base"  # 149M params, 8192-token context, cased

STOPWORDS = {
    "the", "a", "an", "of", "to", "in", "is", "are", "was", "were", "and", "or",
    "for", "on", "with", "as", "by", "that", "this", "it", "be", "could", "you",
    "please", "me", "my", "i", "we", "how", "might", "about", "if", "so", "just",
    "hello", "hi", "hey", "thank", "thanks", "really", "absolutely", "wonderful",
    "great", "sorry",
}
TECH_KEYWORDS = {"python", "json", "api", "sql", "dataframe", "pandas", "openai", "bert"}
NEGATIONS = {
    "not", "no", "never", "none", "cannot", "can't", "don't", "won't",
    "isn't", "aren't", "doesn't", "didn't", "shouldn't", "mustn't", "without",
}

_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
_model = AutoModel.from_pretrained(
    MODEL_ID, attn_implementation="eager", output_attentions=True
).eval()

MAX_TOKENS = 8000      # ModernBERT context is 8192; leave headroom
CHUNK_CHARS = 2000     # ~500 tokens: keeps eager-attention cost near-linear
ATTN_LAYERS = 4        # only the last N layers: cheaper, more semantic signal


def _chunk_text(text):
    """Split over-long text into token-safe chunks at word boundaries."""
    if len(_tokenizer(text)["input_ids"]) <= MAX_TOKENS:
        return [(text, 0)]
    chunks, start = [], 0
    while start < len(text):
        end = min(start + CHUNK_CHARS, len(text))
        if end < len(text):
            space = text.rfind(" ", start, end)
            if space > start:
                end = space
        chunks.append((text[start:end], start))
        start = end + 1 if end < len(text) else end
    return chunks


def _score_chunk(text, offset=0):
    """One forward pass -> (word, attention_importance, sentence_initial, start, end)."""
    enc = _tokenizer(text, return_tensors="pt", return_offsets_mapping=True)
    offsets = enc.pop("offset_mapping")[0]
    word_ids = enc.word_ids()

    with torch.inference_mode():
        outputs = _model(**enc)

    # Mean attention over the last N layers & heads, then attention each
    # token RECEIVES. (All 22 layers on long docs is slow and mostly noise.)
    attn = torch.stack(outputs.attentions[-ATTN_LAYERS:]).mean(dim=(0, 1, 2))
    received = attn.sum(dim=0)

    # Exclude special tokens ([CLS]/[SEP] attention sinks) BEFORE normalizing.
    content_idx = [i for i, w in enumerate(word_ids) if w is not None]
    vals = received[content_idx]
    vals = vals / vals.max()

    # Aggregate subword -> word (max), tracking original char spans.
    word_imp, word_span = {}, {}
    for pos, i in enumerate(content_idx):
        wid = word_ids[i]
        word_imp[wid] = max(word_imp.get(wid, 0.0), vals[pos].item())
        s, e = offsets[i].tolist()
        prev = word_span.get(wid)
        word_span[wid] = (min(prev[0], s), max(prev[1], e)) if prev else (s, e)

    out = []
    for wid in sorted(word_span):
        s, e = word_span[wid]
        prev_char = text[:s].rstrip()[-1:] or ""
        sentence_initial = prev_char in ".!?\n" or s == 0 or not text[:s].strip()
        out.append((text[s:e].strip(), word_imp[wid], sentence_initial,
                    offset + s, offset + e))
    return out


def compress_prompt(text, ratio=0.5, tfidf_weight=0.3):
    # Pass 1: attention importance for every word (chunked if needed).
    words = []
    for chunk, offset in _chunk_text(text):
        words.extend(_score_chunk(chunk, offset))

    # Pass 2: combine with rarity, then apply TIER offsets. Offsets (not
    # multipliers) guarantee ordering: protected > normal > filler, no matter
    # how attention sinks distort the raw scores.
    freq = Counter(w.lower() for w, *_ in words)
    n = max(1, len(words))

    scored = []
    for word, imp, sent_initial, start, end in words:
        lower = word.lower()
        rarity = math.log(1 + n / (1 + freq[lower])) / math.log(1 + n)
        base = (1 - tfidf_weight) * imp + tfidf_weight * rarity

        if not any(ch.isalnum() for ch in word):
            score = base - 100.0     # bare punctuation: drop first
        elif lower in STOPWORDS:
            score = base - 100.0     # filler: kept only to fill the quota
        elif (
            (word[0].isupper() and not sent_initial and lower not in STOPWORDS)
            or lower in TECH_KEYWORDS
            or lower in NEGATIONS                    # never drop "not"/"never"
            or any(ch.isdigit() for ch in word)      # error codes, versions, IDs
        ):
            score = base + 100.0     # protected: survives unless quota is tiny
        else:
            score = base
        scored.append((word, score, start, end))

    n_keep = max(3, int(len(scored) * ratio))
    kept_ids = {id(w) for w, *_ in
                sorted(scored, key=lambda x: x[1], reverse=True)[:n_keep]}
    kept = [(w, s, e) for w, _, s, e in scored if id(w) in kept_ids]

    # Rebuild from the ORIGINAL text: if the gap between two kept words is
    # short punctuation, glue it back -> "created_at", "api.github.com"
    # survive instead of shattering into pieces.
    parts = []
    for i, (word, s, e) in enumerate(kept):
        if i > 0:
            gap = text[kept[i - 1][2]:s]
            # Glue only intra-word punctuation: never weld two sentences together.
            glue = (gap and " " not in gap and "\n" not in gap
                    and len(gap) <= 3 and word[:1].islower()
                    or (gap and " " not in gap and word[:1].isdigit()))
            parts.append(gap if glue else " ")
        parts.append(word)
    return "".join(parts)


# Example Test Run (actual output, ModernBERT-base, ratio=0.5)
prompt = "Could you please write me a Python function that reads a JSON file and converts it into a pandas DataFrame..."
print(compress_prompt(prompt, ratio=0.5))
# -> "write Python function reads JSON file converts into pandas DataFrame"
#    24 tokens -> 13 tokens (-46%), ~65ms on a laptop CPU. Original casing preserved.

BEFORE (32 tokens):
"I was wondering if you could possibly help me understand how I might go about
 resetting the password for my account because I seem to have forgotten it."

AFTER (11 tokens):
"possibly help understand resetting password account because seem"