# Project export: TokenC This document was generated by HackStack to give an AI agent context about a hackathon project. Sections are labeled with their provenance; content marked as truncated was cut to keep this document small. ## Project metadata - Hackathon: UC Berkeley AI Hackathon 2026 - Tagline: Distill Claude into a 66M model that compresses LLM context with~80% with the answer intact. Where classical BM25 drops it, ours keeps it. Fewer tokens, same answers, lower cost. - Devpost: https://devpost.com/software/tokenc - GitHub: https://github.com/gsankar967/tokenc - Team: 0 GitHub contributor(s) — ## Devpost submission (written by the team) ### Inspiration Most compression demos brag about cutting tokens and say nothing about whether the answer survived. Deleting tokens is easy; deleting the right ones is the problem. Right now, frontier LLM models seem to be the best at compressing inputs over classical models like TF-IDF or even lower scale transformer models. Since directly asking frontier LLM like Claude directly is too expensive, we fine tuned a smaller bidirectional encoder based model. Since you can't fine-tune Claude directly, we wanted to see if a small model could learn Claude's sense of what matters and run that compression locally, on every request, for free. ### What it does Given a context and a question, a 66M-parameter model keeps only the sentences needed to answer, down to whatever token budget you set. It's wrapped in an eval that closes the loop (compress, ask Claude, grade the answer) and reports tokens, accuracy, and dollars against a classical baseline, so you see the trade-off instead of trusting it. ### How we built it We framed compression as per-token keep/drop classification instead of generative rewriting. Generation can mangle the exact tokens that matter and costs a decode per request; classification is one forward pass, can't hallucinate, and gives exact label alignment for distillation (no fuzzy-matching a teacher paraphrase back onto the source, which is the noisy step in LLMLingua-2's pipeline). The backbone is a bidirectional encoder (DistilBERT), not a causal SLM like Qwen, because keeping a token depends on context on both sides of it, and a decoder only sees the left. The trainer is model-agnostic (AutoModelForTokenClassification), so the backbone is a one-line swap. For labels, we show Claude the context as a numbered sentence list and ask which indices answer the query. That maps exactly to per-token labels and matches our inference granularity, so train and test agree. Input is Query: … Context: … with query tokens masked out, so relevance is learned conditioned on the question. KEEP is only ~2% of tokens, so we don't threshold at 0.5; we use the model as a ranker (sort by keep-probability) and a separate budget controller fills to the target ratio. That same controller is shared by BM25 and the model, so the comparison changes one variable: the scorer. ### Challenges we ran into BM25 is far stronger than people give it credit for, and understanding why was most of the work. The rare entity name is an exact-match jackpot for a lexical scorer, so on plain retrieval it ties the model. To test semantic compression we built an adversarial slice: synonym queries the answer avoids, plus filler that echoes the query's keyword but carries no answer. Even then it didn't separate until we realized BM25 keeps the whole entity block as a unit, so the answer only gets dropped when the budget is smaller than that block. We tuned the benchmark into that regime. Separately, the class imbalance almost produced an all-DROP model, and our lost-in-the-middle chart came back flat (Haiku reads short contexts fine), so we cut it rather than ship a result that didn't hold. ### Accomplishments we're proud of A model distilled from ~200 Claude-labeled examples that trains in about 10 seconds on a MacBook and produces a real, measured win: at 20% keep (~33% of the tokens) it holds 70% accuracy while BM25 falls to 18%, with full context at 80% and the same budget for both. We're nearly as proud of the restraint: every number comes from Anthropic's own token counter and real API usage, we report where the model only ties classical, and we deleted the chart that lied. ### What we learned Learned compression beats classical only in a specific regime: when relevance is semantic rather than lexical, and when the budget is tight enough to force a real choice. Being able to construct that case on purpose was as much of the work as the model. We also learned how cheap distilling a judgment can be, and that in compression the model is the easy half; proving the answer survived is the hard half, so build the eval first. ### What's next A drop-in proxy: point your Anthropic base_url at TokenC and context gets compressed before it reaches the model, around half off your token bill with one line of config. After that, a stronger backbone and token-level keep/drop for finer control, real long-context benchmarks like HotpotQA and LongBench, and training on diverse multilingual text so it generalizes well past the demo. ## README (from the GitHub repository) # TokenC — distill Claude's context compression into a tiny model **The Token Company Compression Challenge.** Cut the tokens you send an LLM while preserving — and with the hybrid, *improving* — answer quality, by distilling **Claude's relevance judgment** into a small, local **keep/drop token classifier** (the LLMLingua-2 recipe, with Claude as the teacher). ![TokenC performance](performance.png) ### Headline result (from `demo.ipynb`, measured) Semantic multi-doc QA with lexical-trap distractors, downstream reader = **Claude Haiku 4.5**: | approach | accuracy | tokens | |---|---|---| | full context | 80% | 431 | | BM25 (classical) @ 20% keep | 18% | 129 | | trained model (extractive) @ 20% keep | 70% | 141 | | **hybrid (model + rephrase)** | **92%** | **77** | Two results: - The **extractive model stays robust** under aggressive compression — 70% accuracy vs BM25's 18% at the *same* token budget. The rare-token lexical traps fool BM25 into dropping the answer; the model distilled from Claude keeps it. - The **hybrid hits both halves of the challenge at once**: the trained model does the cheap local bulk cut, then a Haiku **densifier** rephrases only the survivors. Result: **92% accuracy at 77 tokens — higher than full context (80%) at ~18% of the tokens.** Denoising the context before the reader sees it removes the distractors that trip it up. The 66M-param compressor trains in **~10 seconds on a MacBook (MPS)**. --- ## Why a learned compressor beats classical BM25 ranks sentences by surface-word overlap. The moment the query and the answer don't share words — *"Where is X **based**?"* answered by *"X **runs everything out of** Helsinki"* — and the context is full of lexical traps (*"X **based** its culture on remote-first work"*), BM25 spends its budget on the traps and **drops the answer**. The model, trained to imitate which sentences **Claude** says are needed, learns the meaning-level mapping and keeps the answer. ## Extractive vs hybrid The trained model is **extractive** — it scores each sentence's keep-probability (query-conditioned) and deletes to a budget. It can only delete, never reword, which is a feature: one cheap forward pass, no hallucination, exact values preserved. The **hybrid** adds an abstractive last mile: after the extractive bulk cut, a Haiku densifier rewrites the surviving sentences into dense facts (values kept verbatim). It runs on a small, already-relevant input, so it's cheap and has little room to invent — and the cost math strongly favors it when the downstream reader is a larger model than the densifier. --- ## How it works ``` (context, query) │ distill.py ▼ Claude (teacher) picks the sentences needed to answer the query → per-token KEEP/DROP labels (data/*.jsonl) │ train_compressor.py ▼ DistilBERT (student) fine-tuned as a query-aware keep/drop classifier → compressor_model/ (class-weighted loss; KEEP is ~2%) │ neural.py ▼ rank sentences by mean KEEP-probability, fill to a token budget → extractive compress(); compress_hybrid() adds a densify() last mile │ tokenc.py (eval harness) ▼ full vs BM25 vs model vs hybrid → accuracy / tokens / $ → demo.ipynb ``` - **Teacher labeling** (`distill.py`): Claude sees the context as a numbered sentence list and returns the indices needed — exact token-label alignment, no fuzzy matching. Mixed lexical + semantic slices; seeds disjoint from the eval set. - **Student** (`train_compressor.py`): `AutoModelForTokenClassification` (default `distilbert-base-uncased`; swap with `--backbone`). A bidirectional encoder is the right architecture for keep/drop — each token sees both directions. Class-weighted loss because KEEP is rare; we treat the model as a ranker, not a 0.5-threshold classifier. - **Inference** (`neural.py`): per-sentence KEEP-probability → rank → fill to budget → re-emit in order. `compress_hybrid()` adds the densify pass. Sentence scores are ratio-independent and disk-cached, so a whole Pareto sweep reuses them. - **Eval** (`tokenc.py`): controllable multi-doc QA with distractors; Claude as the downstream reader; objective substring grading; every LLM call disk-cached so the demo re-runs instantly. --- ## Files | file | what | |---|---| | `tokenc.py` | engine: BM25 baseline, budget controller, benchmark generator, eval harness, `densify`, pricing, caching | | `distill.py` | generate KEEP/DROP labels from Claude (`--offline` for an API-free dry run) | | `train_compressor.py` | fine-tune the keep/drop classifier (`--smoke` for a fast pipeline+timing check) | | `neural.py` | load the trained model; `compress` (extractive) and `compress_hybrid` (extractive + rephrase) | | `make_figures.py` | render `performance.png` and `pareto.png` | | `build_notebook.py` | regenerate `demo.ipynb` | | `demo.ipynb` | the story + charts + interactive keep-rate slider (pre-executed) | | `smoke_test.py` | offline sanity checks (no API key) | --- ## Quickstart ```bash bash setup.sh # venv + deps + Jupyter kernel echo 'ANTHROPIC_API_KEY=sk-ant-...' > .env # auto-loaded by every script ./.venv/bin/python distill.py --n 240 # distill labels from Claude (~5 min, cached) ./.venv/bin/python train_compressor.py # train the model (~10 s on MPS) ./.venv/bin/python make_figures.py # render performance.png / pareto.png ./.venv/bin/jupyter notebook demo.ipynb # open the demo ``` The notebook ships **pre-executed** — open it to see the charts immediately; re-run is instant (cached). The **keep-rate slider** cell is the booth demo: drag it and watch tokens & cost fall while the answer stays correct, then flip to BM25 to watch it break. --- ## Product framing A **drop-in proxy**: point your Anthropic `base_url` at TokenC; we compress every prompt's context before it hits the model — fewer input tokens, same or better answers, one line of config. The extractive compressor is a small model you run locally; the hybrid adds a cheap rephrase pass when you want the last mile. --- ## Honest notes - The benchmark is **synthetic and controllable** by design — it lets us dial the exact regime (lexical traps, compression budget) where classical methods fail and a learned one wins, with objective grading. The methods are content-agnostic and run on any text (try the slider on your own). - The extractive model *maintains* quality under aggressive compression; the **hybrid measurably *improves* over full-context accuracy here (92% vs 80%)** by denoising before the reader. The rephrase is a small extra LLM call on already-pruned input — the cost trade-off favors it when the reader is a bigger model than the densifier. - Models/pricing from the current Claude lineup (Haiku 4.5 $1/$5, Sonnet 4.6 $3/$15, Opus 4.8 $5/$25 per 1M in/out). Token counts use Anthropic's own counter / real API `usage`, never tiktoken. ## Detected evidence (automated analysis) Indexed codebase: 9 recognized source files, 70 KB. - Anthropic (technology) — detected in the code - Python (language) — detected in the code ## Codebase structure (from repository index) ### Files (12 of 12) ``` .gitignore build_notebook.py demo.ipynb distill.py make_figures.py neural.py README.md requirements.txt setup.sh smoke_test.py tokenc.py train_compressor.py ``` ### Dependencies - requirements.txt: anthropic@>=0.69, ipykernel, jupyter, matplotlib, numpy ### Recent commits (newest first) - Show hybrid as a single labeled point on pareto.png - Plot hybrid as a full budget sweep (cluster) on pareto.png, with annotation - Label the hybrid point on pareto.png so it's unmistakable - Add hybrid compressor (extractive + abstractive last mile) and performance figures - TokenC: distill Claude's context compression into a tiny keep/drop model ## Key source files (fetched from GitHub, selected and truncated for size) ### requirements.txt ``` anthropic>=0.69 numpy matplotlib jupyter ipykernel ``` ### setup.sh ```shell #!/usr/bin/env bash # One-shot environment setup for TokenC. set -euo pipefail cd "$(dirname "$0")" python3 -m venv .venv ./.venv/bin/python -m pip install -q --upgrade pip ./.venv/bin/python -m pip install -q -r requirements.txt ./.venv/bin/python -m ipykernel install --user --name tokenc --display-name "TokenC (.venv)" echo echo "Setup done. Next:" echo " 1) echo 'ANTHROPIC_API_KEY=sk-ant-...' > .env" echo " 2) ./.venv/bin/python distill.py --n 240 # distill labels from Claude" echo " 3) ./.venv/bin/python train_compressor.py # train the keep/drop model" echo " 4) ./.venv/bin/python build_notebook.py && ./.venv/bin/jupyter notebook demo.ipynb" ``` ### smoke_test.py ```python """Offline sanity checks for the compression engine — no API key needed. Run: ./.venv/bin/python smoke_test.py Verifies: sentence splitting, BM25 ranking, the budget controller, and that query-aware extraction actually keeps the gold sentence while dropping distractors on a synthetic example. """ import tokenc as tc def test_compress_keeps_answer_drops_distractors(): bench = tc.make_benchmark(n_examples=20, n_docs=10, seed=3) kept_hits = 0 ratios = [] for ex in bench: c = tc.compress(ex.context, ex.question, target_ratio=0.3) ratios.append(c.ratio) # gold value should survive aggressive compression if tc._norm(ex.gold) in tc._norm(c.text): kept_hits += 1 avg_ratio = sum(ratios) / len(ratios) keep_rate = kept_hits / len(bench) print(f"avg kept-ratio at target 0.30 : {avg_ratio:.2f}") print(f"gold-survival rate : {keep_rate*100:.0f}%") assert avg_ratio < 0.45, "compression should hit roughly the target budget" assert keep_rate >= 0.85, "query-aware extraction should keep the gold fact" def test_token_estimate_monotonic(): a = tc.estimate_tokens("hello world") b = tc.estimate_tokens("hello world " * 50) assert b > a > 0 def test_pricing_present(): for m in (tc.DOWNSTREAM_MODEL, "claude-sonnet-4-6", "claude-opus-4-8"): assert m in tc.PRICING if __name__ == "__main__": test_token_estimate_monotonic() test_pricing_present() test_compress_keeps_answer_drops_distractors() print("\nAll offline smoke tests passed ✅") ``` ### make_figures.py ```python """Render submission-ready performance figures from the (cached) eval. Outputs: performance.png — head-to-head: full vs BM25 vs trained model vs hybrid pareto.png — tokens vs accuracy frontier Run: ./.venv/bin/python make_figures.py """ import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import anthropic import tokenc as tc from neural import NeuralCompressor client = anthropic.Anthropic() nc = NeuralCompressor("compressor_model") EVAL = tc.make_benchmark(n_examples=40, n_docs=3, n_filler=8, seed=7, mode="semantic") RATIOS = [1.0, 0.6, 0.5, 0.4, 0.3, 0.2] M = tc.DOWNSTREAM_MODEL bm = tc.run_ratio_sweep(client, EVAL, RATIOS, strategy="extractive", model=M) nz = tc.run_ratio_sweep(client, EVAL, RATIOS, model=M, compress_fn=lambda c, q, r: nc.compress(c, q, r).text) hy = tc.run_ratio_sweep(client, EVAL, [0.15], model=M, compress_fn=lambda c, q, r: nc.compress_hybrid(client, c, q, r).text)[0] full_acc, full_tok = bm[0].accuracy * 100, bm[0].avg_in_tokens def at(res, t): return min(res, key=lambda r: abs(r.ratio_target - t)) bm20, nz20 = at(bm, 0.2), at(nz, 0.2) # ---------------------------------------------------------------- figure 1: bars labels = ["full\ncontext", "BM25\n(classical)", "trained\nmodel\n(extractive)", "hybrid\n(model + rephrase)"] accs = [full_acc, bm20.accuracy * 100, nz20.accuracy * 100, hy.accuracy * 100] toks = [full_tok, bm20.avg_in_tokens, nz20.avg_in_tokens, hy.avg_in_tokens] colors = ["#888888", "#cc4444", "#2277aa", "#22aa88"] fig, ax = plt.subplots(figsize=(8.5, 5.5)) bars = ax.bar(labels, accs, color=colors, width=0.66) for b, a, t in zip(bars, accs, toks): ax.text(b.get_x() + b.get_width() / 2, a + 1.5, f"{a:.0f}%\n{t:.0f} tok", ha="center", va="bottom", fontsize=11, fontweight="bold") ax.axhline(full_acc, ls="--", c="#888888", alpha=.7) ax.text(3.45, full_acc + 0.5, "full-context accuracy", ha="right", va="bottom", fontsize=9, color="#666666") ax.set_ylabel("downstream answer accuracy (%)", fontsize=12) ax.set_ylim(0, 108) ax.set_title("TokenC: accuracy vs tokens under aggressive context compression\n" "semantic multi-doc QA with lexical traps · reader = Claude Haiku 4.5", fontsize=12.5, fontweight="bold") ax.grid(axis="y", alpha=.3) ax.spines[["top", "right"]].set_visible(False) plt.tight_layout() plt.savefig("performance.png", dpi=200, bbox_inches="tight") print("wrote performance.png") # ------------------------------------------------------------- figure 2: pareto fig, ax = plt.subplots(figsize=(8, 5.5)) ax.plot([r.avg_in_tokens for r in bm], [r.accuracy * 100 for r in bm], "o-", color="#cc4444", label="BM25 (classical)") ax.plot([r.avg_in_tokens for r in nz], [r.accuracy * 100 for r in nz], "s-", color="#2277aa", label="trained model (extractive)") ax.scatter([hy.avg_in_tokens], [hy.accuracy * 100], marker="*", s=520, color="#22aa88", edgecolors="#0b5", linewidths=1.2, zorder=6, label="hybrid (model + rephrase)") ax.annotate(f"hybrid\n{hy.accuracy*100:.0f}% @ {hy.avg_in_tokens:.0f} tok", (hy.avg_in_tokens, hy.accuracy * 100), textcoords="offset points", xytext=(16, -2), ha="left", va="center", fontsize=10.5, color="#137a57", fontweight="bold") ax.scatter([full_tok], [full_acc], color="#000000", zorder=5) ax.annotate("full context", (full_tok, full_acc), textcoords="offset points", xytext=(-8, 8), ha="right", fontsize=10) ax.axhline(full_acc, ls="--", c="#888888", alpha=.5) ax.set_xlabel("avg input tokens per request (lower = cheaper)", fontsize=12) ax.set_ylabel("downstream answer accuracy (%)", fontsize=12) ax.set_title("TokenC: token–accuracy frontier\n" "the hybrid beats full-context accuracy at ~18% of the tokens", fontsize=12.5, fontweight="bold") ax.legend(fontsize=10, loc="lower right") ax.grid(alpha=.3) ax.spines[["top", "right"]].set_visible(False) plt.tight_layout() plt.savefig("pareto.png", dpi=200, bbox_inches="tight") print("wrote pareto.png") print(f"\nfull : {full_acc:.0f}% @ {full_tok:.0f} tok") print(f"BM25 : {bm20.accuracy*100:.0f}% @ {bm20.avg_in_tokens:.0f} tok") print(f"model : {nz20.accuracy*100:.0f}% @ {nz20.avg_in_tokens:.0f} tok") print(f"hybrid: {hy.accuracy*100:.0f}% @ {hy.avg_in_tokens:.0f} tok " f"({hy.avg_in_tokens/full_tok*100:.0f}% of full, {hy.accuracy*100-full_acc:+.0f} pts vs full)") ``` ### neural.py ```python """ Inference for the trained keep/drop compressor. Loads the fine-tuned token classifier and turns it into a drop-in compressor with the same interface as the BM25 baseline (`tokenc.compress`). It scores each sentence by the model's mean KEEP-probability (query-conditioned), then the same budget controller keeps top sentences to a target token ratio and re-emits them in original order. Sentence scores are independent of the target ratio, so we compute them once per (context, query) and cache to disk — every ratio in the Pareto sweep reuses them. """ from __future__ import annotations import hashlib import torch import torch.nn.functional as F from transformers import AutoModelForTokenClassification, AutoTokenizer import tokenc as tc def _device(): if torch.backends.mps.is_available(): return "mps" if torch.cuda.is_available(): return "cuda" return "cpu" class NeuralCompressor: def __init__(self, model_dir="compressor_model", device=None, max_len=320, batch_size=16): self.model_dir = str(model_dir) self.device = device or _device() self.tok = AutoTokenizer.from_pretrained(self.model_dir) if self.tok.pad_token is None: self.tok.pad_token = self.tok.eos_token self.model = ( AutoModelForTokenClassification.from_pretrained(self.model_dir) .to(self.device).eval() ) self.max_len = max_len self.batch_size = batch_size # KEEP class index (robust to label order) self.keep_id = self.model.config.label2id.get("KEEP", 1) self._tag = hashlib.sha256(self.model_dir.encode()).hexdigest()[:8] @torch.no_grad() def _score_batch(self, sentences, query): prefix = ["Query:"] + query.split() + ["Context:"] plen = len(prefix) batch_words = [prefix + s.split() for s in sentences] enc = self.tok( batch_words, is_split_into_words=True, truncation=True, max_length=self.max_len, padding=True, return_tensors="pt", ).to(self.device) logits = self.model(**enc).logits keep_prob = F.softmax(logits, dim=-1)[..., self.keep_id] # (B, T) scores = [] for i in range(len(sentences)): wids = enc.word_ids(batch_index=i) vals = [keep_prob[i, t].item() for t, w in enumerate(wids) if w is not None and w >= plen] scores.append(sum(vals) / len(vals) if vals else 0.0) return scores def score_sentences(self, context, query): sents = tc.split_sentences(context) key = f"neural::{self._tag}::{query}::{hashlib.sha256(context.encode()).hexdigest()}" hit = tc._CACHE.get(key) if hit is not None and len(hit["scores"]) == len(sents): return sents, hit["scores"] scores = [] for i in range(0, len(sents), self.batch_size): scores.extend(self._score_batch(sents[i:i + self.batch_size], query)) tc._CACHE.set(key, {"scores": scores}) return sents, scores def compress(self, context, query, target_ratio=0.5, token_fn=tc.estimate_tokens): sents, scores = self.score_sentences(context, query) orig = token_fn(context) if not sents: return tc.Compressed(context, query, "neural", orig, orig, 0, 0) budget = max(1, int(round(orig * target_ratio))) order = sorted(range(len(sents)), key=lambda i: scores[i], reverse=True) kept, used = set(), 0 for i in order: t = token_fn(sents[i]) if kept and used + t > budget: continue kept.add(i); used += t if used >= budget: break text = " ".join(sents[i] for i in sorted(kept)) return tc.Compressed(text, query, "neural", orig, token_fn(text), len(sents), len(kept)) def compress_hybrid(self, client, context, query, target_ratio=0.2, extract_ratio=None, model=tc.COMPRESSOR_MODEL, token_fn=tc.estimate_tokens): """Hybrid: the trained model does the cheap local bulk cut (extractive, no hallucination), then a small LLM rephrases only the surviving text into dense facts for the last mile. The rephrase runs on a small, already-relevant input — so it's cheap and has little room to hallucinate.""" orig = token_fn(context) if extract_ratio is None: extract_ratio = min(1.0, target_ratio * 2.5) pruned = self.compress(context, query, extract_ratio, token_fn) budget = max(20, int(round(orig * target_ratio))) text = tc.densify(client, pruned.text, query, budget, model) return tc.Compressed(text, query, "hybrid", orig, token_fn(text), pruned.n_units_total, pruned.n_units_kept) ``` ### distill.py ```python """ Distill Claude's compression judgment into per-token keep/drop labels. For each (context, query) we show Claude the context as a numbered list of sentences and ask which sentences are *needed* to answer the query. Those become the keep=1 sentences; everything else is drop=0. We emit word-level labels (propagated to sub-word tokens at train time) — this is the LLMLingua-2 setup, with Claude as the teacher. Output: data/distill_train.jsonl, data/distill_val.jsonl Each line: {"query": str, "words": [str, ...], "labels": [0/1, ...]} Run: export ANTHROPIC_API_KEY=sk-... ./.venv/bin/python distill.py --n 240 --teacher claude-haiku-4-5 """ from __future__ import annotations import argparse import json import os import random import re from pathlib import Path import anthropic import tokenc as tc DATA_DIR = Path(__file__).resolve().parent / "data" TEACHER_DEFAULT = "claude-haiku-4-5" # cheap+fast; pass --teacher claude-sonnet-4-6 for sharper labels _TEACHER_SYS = ( "You are a context compressor. You are given a QUERY and a numbered list of " "SENTENCES. Choose the minimal set of sentences strictly required to answer " "the QUERY. Return ONLY a JSON array of integer indices (e.g. [2,5]); no prose." ) def teacher_select(client, sentences: list[str], query: str, model: str) -> set[int]: listing = "\n".join(f"{i}: {s}" for i, s in enumerate(sentences)) user = f"QUERY: {query}\n\nSENTENCES:\n{listing}" key = f"teacher::{model}::{_TEACHER_SYS}::{user}" hit = tc._CACHE.get(key) if hit is not None: raw = hit["raw"] else: resp = client.messages.create( model=model, max_tokens=200, system=_TEACHER_SYS, messages=[{"role": "user", "content": user}], ) raw = "".join(b.text for b in resp.content if b.type == "text").strip() tc._CACHE.set(key, {"raw": raw}) m = re.search(r"\[[\d,\s]*\]", raw) if not m: return set() try: idxs = json.loads(m.group(0)) except json.JSONDecodeError: return set() return {int(i) for i in idxs if 0 <= int(i) < len(sentences)} def build_examples(n: int, seed: int): """Varied (context, query) pairs for training — mixes lexical & semantic modes and doc counts, with seeds far from the eval seeds so we never train on test items. The semantic half is what teaches the student to beat BM25.""" r = random.Random(seed) out, i = [], 0 while len(out) < n: nd = r.choice([4, 6, 8, 10]) nf = r.choice([2, 3, 4, 5]) mode = "semantic" if i % 2 == 0 else "lexical" bench = tc.make_benchmark(n_examples=8, n_docs=nd, n_filler=nf, seed=1000 + i, mode=mode) for ex in bench: out.append(ex) if len(out) >= n: break i += 1 return out def pseudo_select(sentences: list[str], query: str, coverage: float = 0.4) -> set[int]: """Offline pseudo-teacher (BM25 top-k to a coverage budget). Lets us validate the training pipeline end-to-end with no API key. NOT used for the real run.""" tokd = [tc.tokenize(s) for s in sentences] bm25 = tc.BM25(tokd) q = tc.tokenize(query) order = sorted(range(len(sentences)), key=lambda i: bm25.score(q, i), reverse=True) budget = max(1, int(round(sum(tc.estimate_tokens(s) for s in sentences) * coverage))) kept, used = set(), 0 for i in order: t = tc.estimate_tokens(sentences[i]) if kept and used + t > budget: continue kept.add(i); used += t if used >= budget: break return kept def to_labeled_row(ex, kept_idx: set[int]): sents = tc.split_sentences(ex.context) words, labels = [], [] for i, s in enumerate(sents): keep = 1 if i in kept_idx else 0 for w in s.split(): words.append(w) labels.append(keep) return {"query": ex.question, "words": words, "labels": labels} def main(): ap = argparse.ArgumentParser() ap.add_argument("--n", type=int, default=240, help="number of training examples") ap.add_argument("--teacher", default=TEACHER_DEFAULT) ap.add_argument("--seed", type=int, default=2024) ap.add_argument("--val-frac", type=float, default=0.15) ap.add_argument("--offline", action="store_true", help="use BM25 pseudo-labels (no API) to validate the pipeline") args = ap.parse_args() client = None if not args.offline: if not os.environ.get("ANTHROPIC_API_KEY"): raise SystemExit("Set ANTHROPIC_API_KEY first: export ANTHROPIC_API_KEY=sk-... " "(or pass --offline to dry-run with BM25 pseudo-labels)") client = anthropic.Anthropic() DATA_DIR.mkdir(exist_ok=True) examples = build_examples(args.n, args.seed) rows, kept_tot, tok_tot = [], 0, 0 for j, ex in enumerate(examples): sents = tc.split_sentences(ex.context) if args.offline: kept = pseudo_select(sents, ex.question) else: kept = teacher_select(client, sents, ex.question, args.teacher) # Safety net: the teacher should keep the sentence holding the gold value. for i, s in enumerate(sents): if tc._norm(ex.gold) in tc._norm(s): kept.add(i) row = to_labeled_row(ex, kept) rows.append(row) kept_tot += sum(row["labels"]) tok_tot += len(row["labels"]) if (j + 1) % 20 == 0: print(f" labeled {j+1}/{len(examples)} " f"(keep-rate so far {kept_tot/max(1,tok_tot)*100:.0f}%)") random.Random(args.seed).shuffle(rows) n_val = max(1, int(len(rows) * args.val_frac)) val, train = rows[:n_val], rows[n_val:] (DATA_DIR / "distill_train.jsonl").write_text( "\n".join(json.dumps(r) for r in train)) (DATA_DIR / "distill_val.jsonl").write_text( "\n".join(json.dumps(r) for r in val)) [truncated — 259 more characters] ``` ### train_compressor.py ```python """ Train the query-aware keep/drop token classifier (LLMLingua-2 recipe). Model-agnostic via AutoModelForTokenClassification: * default backbone: distilbert-base-uncased (small bidirectional encoder — the right architecture for keep/drop; each token sees both directions) * stronger options: bert-base-uncased, microsoft/deberta-v3-small, answerdotai/ModernBERT-base (just pass --backbone) Input format per example: "Query: Context: " - query/prefix tokens get label -100 (ignored in loss) - context word tokens get the keep/drop label (first sub-token labeled) Run: # quick pipeline+timing check (offline pseudo-labels) ./.venv/bin/python distill.py --offline --n 80 ./.venv/bin/python train_compressor.py --smoke # real run after distilling from Claude ./.venv/bin/python train_compressor.py --backbone distilbert-base-uncased --epochs 3 """ from __future__ import annotations import argparse import json import time from pathlib import Path import numpy as np import torch import torch.nn.functional as F from datasets import Dataset from transformers import ( AutoModelForTokenClassification, AutoTokenizer, DataCollatorForTokenClassification, Trainer, TrainingArguments, ) DATA_DIR = Path(__file__).resolve().parent / "data" OUT_DEFAULT = Path(__file__).resolve().parent / "compressor_model" def load_rows(path: Path): return [json.loads(l) for l in path.read_text().splitlines() if l.strip()] def make_dataset(rows, tokenizer, max_len: int): def encode(row): prefix = ["Query:"] + row["query"].split() + ["Context:"] words = prefix + row["words"] wlabels = [-100] * len(prefix) + row["labels"] enc = tokenizer( words, is_split_into_words=True, truncation=True, max_length=max_len, ) word_ids = enc.word_ids() labels, prev = [], None for wid in word_ids: if wid is None: labels.append(-100) elif wid != prev: labels.append(wlabels[wid]) # label first sub-token of a word else: labels.append(-100) # ignore continuation sub-tokens prev = wid enc["labels"] = labels return enc return Dataset.from_list([encode(r) for r in rows]) def keep_f1(pred): logits, labels = pred preds = np.argmax(logits, axis=-1) mask = labels != -100 p, y = preds[mask], labels[mask] tp = int(((p == 1) & (y == 1)).sum()) fp = int(((p == 1) & (y == 0)).sum()) fn = int(((p == 0) & (y == 1)).sum()) prec = tp / (tp + fp) if tp + fp else 0.0 rec = tp / (tp + fn) if tp + fn else 0.0 f1 = 2 * prec * rec / (prec + rec) if prec + rec else 0.0 acc = float((p == y).mean()) if len(y) else 0.0 return {"keep_precision": prec, "keep_recall": rec, "keep_f1": f1, "token_acc": acc} class WeightedTrainer(Trainer): """Token-classification with a class-weighted loss — the KEEP class is rare (~2%, since usually one sentence answers the query), so we up-weight it.""" def __init__(self, *args, class_weights=None, **kwargs): super().__init__(*args, **kwargs) self._class_weights = class_weights def compute_loss(self, model, inputs, return_outputs=False, **kwargs): labels = inputs.pop("labels") outputs = model(**inputs) logits = outputs.logits w = None if self._class_weights is None else self._class_weights.to(logits.device) loss = F.cross_entropy( logits.view(-1, logits.size(-1)), labels.view(-1), weight=w, ignore_index=-100, ) return (loss, outputs) if return_outputs else loss def main(): ap = argparse.ArgumentParser() ap.add_argument("--backbone", default="distilbert-base-uncased") ap.add_argument("--keep-weight", type=float, default=8.0, help="loss weight for the rare KEEP class") ap.add_argument("--epochs", type=float, default=3.0) ap.add_argument("--bs", type=int, default=8) ap.add_argument("--lr", type=float, default=2e-5) ap.add_argument("--max-len", type=int, default=320) ap.add_argument("--out", default=str(OUT_DEFAULT)) ap.add_argument("--smoke", action="store_true", help="1 epoch on a tiny subset to validate pipeline + timing") args = ap.parse_args() device = ("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu") print(f"Backbone: {args.backbone} | device: {device}") train_rows = load_rows(DATA_DIR / "distill_train.jsonl") val_rows = load_rows(DATA_DIR / "distill_val.jsonl") if args.smoke: train_rows, val_rows = train_rows[:16], val_rows[:8] args.epochs = 1.0 tok = AutoTokenizer.from_pretrained(args.backbone) if tok.pad_token is None: # Qwen et al. have no pad token tok.pad_token = tok.eos_token model = AutoModelForTokenClassification.from_pretrained( args.backbone, num_labels=2, id2label={0: "DROP", 1: "KEEP"}, label2id={"DROP": 0, "KEEP": 1}, ) if model.config.pad_token_id is None: model.config.pad_token_id = tok.pad_token_id ds_train = make_dataset(train_rows, tok, args.max_len) ds_val = make_dataset(val_rows, tok, args.max_len) collator = DataCollatorForTokenClassification(tok) targs = TrainingArguments( output_dir=str(Path(args.out) / "_trainer"), num_train_epochs=args.epochs, per_device_train_batch_size=args.bs, per_device_eval_batch_size=args.bs, learning_rate=args.lr, eval_strategy="epoch", save_strategy="no", logging_steps=10, report_to=[], fp16=False, bf16=False, ) trainer = WeightedTrainer( model=model, args=targs, train_dataset=ds_train, eval_dataset=ds_val, data_collator=collator, compute_metrics=keep_f1, class_w [truncated — 1385 more characters] ``` ### build_notebook.py ```python """Builds demo.ipynb from cell sources (safer than hand-writing JSON). Run: ./.venv/bin/python build_notebook.py """ import nbformat as nbf C = [] # (kind, source) def md(s): C.append(("md", s)) def code(s): C.append(("code", s)) md(r"""# TokenC — distill Claude's context compression into a tiny model **Thesis (The Token Company):** cut the tokens you send an LLM by ~50% while *preserving or improving* answer quality. **What this notebook shows, with measurements (not vibes):** 1. **A trained compressor** — a small bidirectional encoder fine-tuned as a query-aware *keep/drop* token classifier (the LLMLingua-2 recipe), distilled from **Claude** as the teacher. 2. **Pareto curve** — tokens vs downstream accuracy: full context vs **BM25 (classical baseline)** vs **our trained model**. The learned model holds accuracy at far fewer tokens — *classical methods can't, once the query and answer don't share words.* 3. **The money** — $ saved per 1M requests at real Claude prices. 4. **A hybrid that cuts tokens *and* improves accuracy** — the trained model does the bulk cut, a Haiku densifier rephrases the survivors, and the result beats full-context accuracy at ~18% of the tokens. 5. An interactive **keep-rate slider** for the booth. **Run order:** `distill.py` → `train_compressor.py` → this notebook. If the trained model isn't present yet, the neural cells fall back to BM25 so everything still runs.""") code(r"""import os, numpy as np, matplotlib.pyplot as plt import anthropic, tokenc as tc %matplotlib inline try: from neural import NeuralCompressor except Exception as _e: NeuralCompressor = None print("neural import failed:", _e) assert os.environ.get("ANTHROPIC_API_KEY"), "No ANTHROPIC_API_KEY — put it in .env" client = anthropic.Anthropic() DOWNSTREAM = tc.DOWNSTREAM_MODEL # the 'reader' the compressor feeds RATIOS = [1.0, 0.6, 0.5, 0.4, 0.3, 0.2] # 1.0 = full context (no compression) neural = None if NeuralCompressor and os.path.isdir("compressor_model"): neural = NeuralCompressor("compressor_model") print("Loaded trained neural compressor:", neural.model_dir) else: print("No trained model yet -> neural cells fall back to BM25.") print("Downstream reader:", DOWNSTREAM)""") md(r"""## The method For each `(context, query)` we ask **Claude** which sentences are actually needed to answer the query. Those become per-token **keep/drop** labels, and we fine-tune a small encoder to imitate that judgment — a **0.5%-of-Claude-size model that compresses with Claude's relevance sense**. At inference it scores each sentence by KEEP-probability (query-conditioned) and a budget controller keeps the top sentences to a target token ratio. **BM25** (lexical sentence ranking) is the classical baseline we measure against.""") code(r"""# Live compression on one hard (semantic, trap-laden) example — real Claude token counts. ex = tc.make_benchmark(n_examples=1, n_docs=4, n_filler=8, seed=4242, mode="semantic")[0] print("Question :", ex.question) print("Gold :", ex.gold, "\n") full_tok = tc.count_tokens(client, ex.context) print("FULL context tokens (Claude counter):", full_tok, "\n") bm = tc.compress(ex.context, ex.question, target_ratio=0.25) print("BM25 ->", bm.summary()) if neural: nz = neural.compress(ex.context, ex.question, target_ratio=0.25) print("MODEL ->", nz.summary()) print("\nModel kept:\n", nz.text) else: print("\nBM25 kept:\n", bm.text)""") md(r"""## Keep-rate slider (interactive booth demo) Drag the **keep rate** down and watch tokens & cost fall while the answer stays correct — then flip to the lexical-only compressor to see it break where the learned model holds.""") code(r"""import ipywidgets as W from IPython.display import display, clear_output, Markdown S = tc.make_benchmark(n_examples=1, n_docs=4, n_filler=8, seed=77, mode="semantic")[0] ORIG = tc.ask(client, S.context, S.question, DOWNSTREAM) # full baseline (cached) price = tc.PRICING[DOWNSTREAM]["in"] slider = W.FloatSlider(value=0.30, min=0.10, max=1.0, step=0.05, description="keep rate", continuous_update=False, readout_format=".0%") which = W.ToggleButtons(options=(["BM25"] + (["Trained model"] if neural else [])), description="compressor") out = W.Output() def render(*_): with out: clear_output() r = slider.value if r >= 0.999: text = S.context elif which.value == "Trained model" and neural: text = neural.compress(S.context, S.question, r).text else: text = tc.compress(S.context, S.question, r).text a = tc.ask(client, text, S.question, DOWNSTREAM) ok = tc.graded_correct(a["answer"], S.gold) saved = (1 - a["in_tokens"] / ORIG["in_tokens"]) * 100 display(Markdown( f"**Q:** {S.question}\n\n" f"**Prompt tokens:** {a['in_tokens']} (full {ORIG['in_tokens']}) - **{saved:.0f}% fewer**\n\n" f"**Input $ / 1M requests:** ${a['in_tokens']*price:,.0f}\n\n" f"**Answer:** `{a['answer']}` -> {'CORRECT' if ok else 'WRONG'} (gold: {S.gold})")) slider.observe(render, "value"); which.observe(render, "value") display(W.VBox([slider, which, out])); render()""") md(r"""## Benchmark — multi-doc QA with distractors (semantic slice) Each question targets one entity's attribute. The answer document sits in the **middle** of distractor documents about other entities. In the **semantic** slice the query uses synonyms and the answer sentence avoids the query's words — so lexical overlap no longer identifies the answer.""") code(r"""# Semantic, trap-laden slice: each entity's block of lexical-trap fillers exceeds # the token budget, so BM25 spends the budget on traps and drops the real answer. EVAL = tc.make_benchmark(n_examples=40, n_docs=3, n_filler=8, seed=7, mode="semantic") avg = round(np.mean([tc.estimate_tokens(e.context) for e in EVAL])) print(len(EVAL), "exam [truncated — 8893 more characters] ``` ### tokenc.py ```python """ TokenC — query-aware context compression for LLMs. The pitch: send fewer tokens to the model while preserving (and often improving) answer quality. This module is the engine + an eval harness that *proves* the claim with a token-reduction-vs-downstream-quality curve on a controllable multi-doc QA benchmark with distractors. Design notes ------------ * The default compressor is **query-aware extractive selection** (a compact BM25 ranker). It is fast, dependency-free, deterministic, and — crucially — removes distractor/irrelevant text, which is exactly what makes the downstream model *more* accurate on long noisy contexts ("lost in the middle"). * An optional **LLM densifier** (Haiku) is included as a second strategy. * All LLM calls are cached to disk so the notebook re-runs instantly and cheaply during a live demo. Nothing here uses tiktoken. Token counts come from Anthropic's own counter / the real `usage` returned by the Messages API. """ from __future__ import annotations import hashlib import json import math import os import random import re import time from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field, asdict from pathlib import Path from typing import Callable, Optional def _pmap(fn, items, workers: int = 8): """Thread-pooled map that preserves order (LLM calls are IO-bound).""" if workers <= 1: return [fn(x) for x in items] with ThreadPoolExecutor(max_workers=workers) as ex: return list(ex.map(fn, items)) # ---------------------------------------------------------------------------- # Models & pricing (USD per 1,000,000 tokens) — current Claude lineup. # ---------------------------------------------------------------------------- PRICING = { "claude-haiku-4-5": {"in": 1.0, "out": 5.0}, "claude-sonnet-4-6": {"in": 3.0, "out": 15.0}, "claude-opus-4-8": {"in": 5.0, "out": 25.0}, } DOWNSTREAM_MODEL = "claude-haiku-4-5" # the "reader" the compressor feeds COMPRESSOR_MODEL = "claude-haiku-4-5" # used only by the LLM densifier strategy CACHE_DIR = Path(__file__).resolve().parent / ".tokenc_cache" def _load_dotenv() -> None: """Minimal .env loader (no dependency). Lets every script + the notebook pick up ANTHROPIC_API_KEY from a local, gitignored .env without hardcoding it.""" if os.environ.get("ANTHROPIC_API_KEY"): return envp = Path(__file__).resolve().parent / ".env" if not envp.exists(): return for line in envp.read_text().splitlines(): line = line.strip() if not line or line.startswith("#") or "=" not in line: continue k, v = line.split("=", 1) os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'")) _load_dotenv() # ---------------------------------------------------------------------------- # Tiny disk cache so a live demo never pays twice for the same call. # ---------------------------------------------------------------------------- class DiskCache: def __init__(self, path: Path = CACHE_DIR): self.path = Path(path) self.path.mkdir(parents=True, exist_ok=True) def _file(self, key: str) -> Path: h = hashlib.sha256(key.encode("utf-8")).hexdigest()[:32] return self.path / f"{h}.json" def get(self, key: str): f = self._file(key) if f.exists(): return json.loads(f.read_text()) return None def set(self, key: str, value) -> None: self._file(key).write_text(json.dumps(value)) _CACHE = DiskCache() # ---------------------------------------------------------------------------- # Token counting. # * estimate_tokens : instant, offline, monotonic — used for the budget knob. # * count_tokens : exact, via Anthropic's counter — used for headline numbers. # ---------------------------------------------------------------------------- def estimate_tokens(text: str) -> int: """Fast offline token estimate. ~chars/4, the standard rough heuristic.""" if not text: return 0 return max(1, round(len(text) / 4)) def count_tokens(client, text: str, model: str = DOWNSTREAM_MODEL) -> int: """Exact prompt token count via Anthropic's token-counting endpoint (cached).""" key = f"count::{model}::{text}" hit = _CACHE.get(key) if hit is not None: return hit["input_tokens"] resp = client.messages.count_tokens( model=model, messages=[{"role": "user", "content": text}] ) _CACHE.set(key, {"input_tokens": resp.input_tokens}) return resp.input_tokens # ---------------------------------------------------------------------------- # Text utilities: sentence splitting + a minimal stemmer for robust matching. # ---------------------------------------------------------------------------- _SENT_SPLIT = re.compile(r"(?<=[.!?])\s+(?=[A-Z0-9])") _WORD = re.compile(r"[A-Za-z0-9]+") _STOP = set( "a an the of to in on at for and or but is are was were be been being " "this that these those with as by from it its their his her our your my " "what which who whom whose when where why how do does did has have had " "will would can could should may might into about over under than then".split() ) def split_sentences(text: str) -> list[str]: """Split into sentence-ish units, also breaking on hard newlines.""" units: list[str] = [] for block in re.split(r"\n{2,}", text.strip()): block = block.strip() if not block: continue parts = _SENT_SPLIT.split(block.replace("\n", " ")) units.extend(p.strip() for p in parts if p.strip()) return units def _stem(tok: str) -> str: for suf in ("ing", "edly", "ed", "ly", "es", "s"): if len(tok) > len(suf) + 2 and tok.endswith(suf): return tok[: -len(suf)] return tok def tokenize(text: str) -> list[str]: return [_stem(w) for w in _WORD.findall(text.lower()) if w not in _STOP] # ----------------------------------- [truncated — 18173 more characters] ```