ctxprune 0.1.0__py3-none-any.whl

ctxprune/__init__.py

@@ -0,0 +1,6 @@
+"""ctxprune: context compression for AI agents."""
+
+from .compress import Compressor
+
+__all__ = ["Compressor"]
+__version__ = "0.1.0"

ctxprune/align.py

@@ -0,0 +1,93 @@
+"""Turn a teacher's deletion-only compression into per-atom keep labels.
+
+The teacher is asked to only delete atoms, so its output should be an
+ordered subsequence of the source atoms. We find the longest common
+subsequence (exact match first, casefold second) and label matched source
+atoms as kept. Teacher atoms that match nothing are "variations" (the
+teacher rewrote or invented text); a high variation rate means the sample
+is untrustworthy and gets filtered.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .atoms import Atom, atomize, is_protected, norm
+
+
+@dataclass
+class Alignment:
+    keep: list[bool]
+    n_src: int
+    n_comp: int
+    n_matched: int
+    variation_rate: float  # teacher atoms with no source match
+    comp_rate: float  # kept / source, in atoms
+    protected_total: int
+    protected_kept: int
+
+
+def _lcs_pairs(a: list[str], b: list[str]) -> list[tuple[int, int]]:
+    n, m = len(a), len(b)
+    if n == 0 or m == 0:
+        return []
+    # dp[i][j] = LCS length of a[i:], b[j:]; rows built bottom-up.
+    dp = [[0] * (m + 1) for _ in range(n + 1)]
+    for i in range(n - 1, -1, -1):
+        ai, row, nxt = a[i], dp[i], dp[i + 1]
+        for j in range(m - 1, -1, -1):
+            if ai == b[j]:
+                row[j] = nxt[j + 1] + 1
+            else:
+                x, y = nxt[j], row[j + 1]
+                row[j] = x if x >= y else y
+    pairs, i, j = [], 0, 0
+    while i < n and j < m:
+        if a[i] == b[j]:
+            pairs.append((i, j))
+            i += 1
+            j += 1
+        elif dp[i + 1][j] >= dp[i][j + 1]:
+            i += 1
+        else:
+            j += 1
+    return pairs
+
+
+def align(src_atoms: list[Atom], comp_text: str) -> Alignment:
+    src = [a.text for a in src_atoms]
+    comp = [a.text for a in atomize(comp_text) if a.text]
+    keep = [False] * len(src)
+    matched_comp = [False] * len(comp)
+
+    # Pass 1: exact. Pass 2: casefold over what is left, in the gaps between
+    # pass-1 anchors so order is preserved.
+    pairs = _lcs_pairs(src, comp)
+    for i, j in pairs:
+        keep[i] = True
+        matched_comp[j] = True
+    anchors = [(-1, -1)] + pairs + [(len(src), len(comp))]
+    for (i0, j0), (i1, j1) in zip(anchors, anchors[1:]):
+        si = [i for i in range(i0 + 1, i1) if not keep[i]]
+        cj = [j for j in range(j0 + 1, j1) if not matched_comp[j]]
+        if not si or not cj:
+            continue
+        sub = _lcs_pairs([norm(src[i]) for i in si], [norm(comp[j]) for j in cj])
+        for a, b in sub:
+            keep[si[a]] = True
+            matched_comp[cj[b]] = True
+
+    real = [i for i, t in enumerate(src) if t]
+    n_src = len(real)
+    n_matched = sum(matched_comp)
+    prot = [i for i in real if is_protected(src[i])]
+    return Alignment(
+        keep=keep,
+        n_src=n_src,
+        n_comp=len(comp),
+        n_matched=n_matched,
+        variation_rate=(1 - n_matched / len(comp)) if comp else 0.0,
+        comp_rate=(sum(keep[i] for i in real) / n_src) if n_src else 0.0,
+        protected_total=len(prot),
+        protected_kept=sum(keep[i] for i in prot),
+    )

ctxprune/atoms.py

@@ -0,0 +1,106 @@
+"""Atom segmentation: the unit the compressor keeps or drops.
+
+An atom is either a "word" that glues identifier-ish characters together
+(paths, URLs, versions, IPs, snake_case, UUIDs) or a single punctuation
+character. Each atom remembers the whitespace that preceded it, so
+``detok(atoms, keep)`` rebuilds kept text exactly as it appeared in the
+source, without the "0. 21" / "ord _ 8f3a" artifacts LLMLingua-2 produces.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+# Scripts written without spaces (Han, kana) get one atom per character, so
+# a deletion inside a sentence is still a deletion and not a "rewrite".
+_CJK = "\u3040-\u30ff\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff"
+_W = rf"[^\W{_CJK}]"  # word char that is not CJK
+# A word may start with one path/flag sigil, then word chars, with
+# identifier punctuation allowed only *between* word chars.
+_WORD = rf"(?:--?|[/~.$@#])?{_W}(?:(?:{_W}|[.:/@%+#~=?&'\-])*{_W})?"
+_ATOM_RE = re.compile(rf"(\s*)({_WORD}|[{_CJK}]|[^\s])", re.UNICODE)
+
+# All-caps English that is prose, not an identifier (license headers, shouting).
+_CAPS_WORDS = set(
+    "THE AND FOR ANY NOT BUT ARE WAS YOU ALL WITH FROM THIS THAT WITHOUT WARRANTIES WARRANTY SOFTWARE KIND "
+    "IMPLIED INCLUDING LIMITED COPYRIGHT OTHER LIABILITY USE BASIS CONDITIONS SUCH EVENT SHALL HOLDERS "
+    "CONTRIBUTORS DAMAGES PURPOSE PARTICULAR FITNESS MERCHANTABILITY EXPRESS OUT CONNECTION ARISING "
+    "WHETHER CONTRACT TORT OTHERWISE PROVIDED LICENSE NOTICE NOTE IMPORTANT".split()
+)
+_HEX_RE = re.compile(r"^[0-9a-fA-F]{7,}$")
+_CAMEL_RE = re.compile(r"[a-z][A-Z]|[A-Z]{2}[a-z]")
+
+
+@dataclass(frozen=True)
+class Atom:
+    ws: str  # whitespace preceding the atom in the source
+    text: str
+
+
+def atomize(text: str) -> list[Atom]:
+    atoms: list[Atom] = []
+    pos = 0
+    for m in _ATOM_RE.finditer(text):
+        if m.start() != pos:  # unreachable by construction; guard anyway
+            raise ValueError(f"atomizer skipped {text[pos:m.start()]!r}")
+        atoms.append(Atom(m.group(1), m.group(2)))
+        pos = m.end()
+    tail = text[pos:]
+    if tail.strip():
+        raise ValueError(f"atomizer left non-space tail {tail!r}")
+    if tail:  # trailing whitespace rides on a sentinel-free empty atom
+        atoms.append(Atom(tail, ""))
+    return atoms
+
+
+def detok(atoms: list[Atom], keep: list[bool], collapse_ws: bool = True) -> str:
+    """Rebuild text from kept atoms.
+
+    Kept atoms keep their own leading whitespace. When an atom's predecessor
+    was dropped, the dropped run's whitespace is collapsed to the "largest"
+    whitespace seen in that run (a newline wins over a space), so line
+    structure survives deletions.
+    """
+    out: list[str] = []
+    pending = ""
+    for a, k in zip(atoms, keep):
+        if k:
+            ws = a.ws
+            if pending and collapse_ws:
+                ws = _merge_ws(pending, a.ws)
+            if not out:
+                ws = ws if not collapse_ws else ""
+            out.append(ws + a.text)
+            pending = ""
+        else:
+            pending = _merge_ws(pending, a.ws) if pending else (a.ws or " ")
+    return "".join(out)
+
+
+def _merge_ws(a: str, b: str) -> str:
+    if "\n" in a or "\n" in b:
+        return "\n"
+    return " " if (a or b) else ""
+
+
+def is_protected(atom: str) -> bool:
+    """Atoms an agent is likely to need verbatim later (IDs, numbers, paths)."""
+    if not atom or not any(c.isalnum() for c in atom):
+        return False
+    if any(c.isdigit() for c in atom):
+        return True
+    core = atom.lstrip("-/~.$@#")
+    if any(c in core for c in "_/.:@="):
+        return True
+    if _HEX_RE.match(core):
+        return True
+    if len(core) >= 3 and core.isupper() and core not in _CAPS_WORDS:
+        return True
+    if _CAMEL_RE.search(core):
+        return True
+    return atom != core  # flags like --format, paths like /tmp
+
+
+def norm(atom: str) -> str:
+    return atom.casefold()

ctxprune/chunking.py

@@ -0,0 +1,104 @@
+"""Split documents into chunks that fit the compressor's 512-token window.
+
+Chunks break on line boundaries where possible; a single line longer than
+the budget is split between atoms. Token counts use the target encoder's
+tokenizer, so every training chunk fits one forward pass, matching how
+llmlingua feeds the model (``max_seq_len = 512`` minus CLS/SEP).
+"""
+
+from __future__ import annotations
+
+import hashlib
+import random
+from collections.abc import Callable
+
+from .atoms import atomize
+
+DEFAULT_TOKENIZER = "jhu-clsp/mmBERT-small"
+
+
+def load_counter(name: str = DEFAULT_TOKENIZER) -> Callable[[list[str]], list[int]]:
+    from transformers import AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(name)
+
+    def count(texts: list[str]) -> list[int]:
+        return [len(ids) for ids in tok(texts, add_special_tokens=False)["input_ids"]]
+
+    return count
+
+
+def _split_long_line(line: str, budget: int, count) -> list[str]:
+    atoms = atomize(line)
+    pieces = [a.ws + a.text for a in atoms]
+    sizes = count(pieces)
+    out, cur, cur_n = [], [], 0
+    for p, n in zip(pieces, sizes):
+        if cur and cur_n + n > budget:
+            out.append("".join(cur))
+            cur, cur_n = [], 0
+        cur.append(p)
+        cur_n += n
+    if cur:
+        out.append("".join(cur))
+    return out
+
+
+def chunk_text(text: str, count, max_tokens: int = 500, min_tokens: int = 48) -> list[str]:
+    lines = text.splitlines(keepends=True)
+    if not lines:
+        return []
+    sizes = count(lines)
+    units: list[tuple[str, int]] = []
+    for line, n in zip(lines, sizes):
+        if n <= max_tokens:
+            units.append((line, n))
+        else:
+            for piece in _split_long_line(line, max_tokens, count):
+                units.append((piece, count([piece])[0]))
+    chunks, cur, cur_n = [], [], 0
+    for u, n in units:
+        if cur and cur_n + n > max_tokens:
+            chunks.append("".join(cur))
+            cur, cur_n = [], 0
+        cur.append(u)
+        cur_n += n
+    if cur:
+        chunks.append("".join(cur))
+    # Line-sum can undercount at joins; re-check and drop tiny tails.
+    final = []
+    for c, n in zip(chunks, count(chunks)):
+        if n < min_tokens or not c.strip():
+            continue
+        if n > max_tokens:
+            final.extend(p for p in _split_long_line(c, max_tokens, count) if p.strip())
+        else:
+            final.append(c)
+    return final
+
+
+def split_of(group: str, val_pct: int = 2, test_pct: int = 2) -> str:
+    b = int(hashlib.sha1(group.encode()).hexdigest(), 16) % 100
+    if b < test_pct:
+        return "test"
+    if b < test_pct + val_pct:
+        return "val"
+    return "train"
+
+
+def chunk_doc(doc: dict, count, rng: random.Random, max_tokens: int = 500,
+              max_chunks_per_doc: int = 2) -> list[dict]:
+    chunks = chunk_text(doc["text"], count, max_tokens=max_tokens)
+    if len(chunks) > max_chunks_per_doc:
+        idx = sorted(rng.sample(range(len(chunks)), max_chunks_per_doc))
+        chunks = [chunks[i] for i in idx]
+    out = []
+    for c in chunks:
+        meta = {k: v for k, v in doc.items() if k != "text"}
+        out.append({
+            "id": f"{doc['source']}:{hashlib.sha1(c.encode()).hexdigest()[:16]}",
+            "text": c,
+            "split": split_of(doc["group"]),
+            **meta,
+        })
+    return out

ctxprune/cli.py

@@ -0,0 +1,40 @@
+"""ctxprune command line: compress a file (or stdin) and print the result.
+
+    ctxprune order.json --rate 0.5
+    cat build.log | ctxprune --rate 0.33 --force-protected
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+DEFAULT_MODEL = "darioooooo0o/ctxprune-small"
+
+
+def main(argv: list[str] | None = None) -> None:
+    ap = argparse.ArgumentParser(prog="ctxprune", description="Compress text for an LLM by deleting low-value words.")
+    ap.add_argument("file", nargs="?", help="input file (default: stdin)")
+    ap.add_argument("--rate", type=float, default=0.5, help="fraction of the text to keep (default 0.5)")
+    ap.add_argument("--force-protected", action="store_true", help="never drop identifiers, numbers or paths")
+    ap.add_argument("--model", default=DEFAULT_MODEL)
+    ap.add_argument("--backend", choices=["onnx", "torch"], default=None,
+                    help="default: onnx if onnxruntime is installed, else torch")
+    args = ap.parse_args(argv)
+
+    text = open(args.file, encoding="utf-8").read() if args.file else sys.stdin.read()
+    backend = args.backend
+    if backend is None:
+        try:
+            import onnxruntime  # noqa: F401
+            backend = "onnx"
+        except ImportError:
+            backend = "torch"
+    from .compress import Compressor
+
+    c = Compressor(args.model, backend=backend)
+    sys.stdout.write(c.compress(text, rate=args.rate, force_protected=args.force_protected)["text"] + "\n")
+
+
+if __name__ == "__main__":
+    main()

ctxprune/compress.py

@@ -0,0 +1,182 @@
+"""Inference: score atoms with the trained token classifier and keep the best.
+
+    from ctxprune.compress import Compressor
+    c = Compressor("path/or/hub-id")
+    c.compress(text, rate=0.5)["text"]
+
+`rate` is the fraction of tokens to keep (like LLMLingua-2's `rate`).
+Identifiers are never split or re-spaced, because selection happens on whole
+atoms and text is rebuilt from the source's own whitespace.
+"""
+
+from __future__ import annotations
+
+import math
+
+from .atoms import atomize, detok, is_protected
+from .tokenlabels import atom_spans, char_to_atom, special_wrap, token_atom_sets, token_atoms
+
+STRUCTURE = set('{}[]",:')
+
+
+class Compressor:
+    def __init__(self, model: str, device: str | None = None, max_len: int = 512, batch_size: int = 16,
+                 backend: str = "torch", onnx_file: str = "onnx/model.onnx"):
+        """backend="torch" (any device) or "onnx" (CPU, no torch needed; onnx_file may be
+        "onnx/model_quantized.onnx" for the int8 build)."""
+        if backend == "onnx":
+            import os
+
+            os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")  # silence "PyTorch was not found"
+        from transformers import AutoTokenizer
+
+        self.tok = AutoTokenizer.from_pretrained(model)
+        self.backend = backend
+        if backend == "onnx":
+            import os
+
+            import onnxruntime as ort
+            from huggingface_hub import hf_hub_download
+
+            path = os.path.join(model, onnx_file) if os.path.isdir(model) else hf_hub_download(model, onnx_file)
+            self.session = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
+            self.device = "cpu"
+        else:
+            import torch
+            from transformers import AutoModelForTokenClassification
+
+            self.torch = torch
+            if device is None:
+                device = "xpu" if hasattr(torch, "xpu") and torch.xpu.is_available() else (
+                    "cuda" if torch.cuda.is_available() else "cpu")
+            self.device = device
+            self.model = AutoModelForTokenClassification.from_pretrained(model).to(device).eval()
+        self.max_len = max_len
+        self.batch_size = batch_size
+        self.pre, self.post = special_wrap(self.tok)
+
+    def atom_scores(self, text: str) -> tuple[list, list[float], list[int]]:
+        """Keep-probability and token count for every atom of `text`."""
+        atoms = atomize(text)
+        full, spans = atom_spans(atoms)
+        owner = char_to_atom(spans, len(full))
+        enc = self.tok(full, add_special_tokens=False, return_offsets_mapping=True)
+        ids, offs = enc["input_ids"], enc["offset_mapping"]
+        tok_atom = token_atoms(offs, owner)
+
+        # Windows of max_len - 2 tokens, cut where a new atom starts.
+        win = self.max_len - len(self.pre) - len(self.post)
+        windows, start = [], 0
+        while start < len(ids):
+            end = min(start + win, len(ids))
+            if end < len(ids):
+                cut = end
+                while cut > start + 1 and tok_atom[cut] == tok_atom[cut - 1]:
+                    cut -= 1
+                end = cut if cut > start + 1 else end
+            windows.append((start, end))
+            start = end
+
+        probs = [0.0] * len(ids)
+        for b in range(0, len(windows), self.batch_size):
+            batch = windows[b:b + self.batch_size]
+            seqs = [self.pre + ids[s:e] + self.post for s, e in batch]
+            width = max(map(len, seqs))
+            pad = self.tok.pad_token_id
+            ids_rows = [q + [pad] * (width - len(q)) for q in seqs]
+            mask_rows = [[1] * len(q) + [0] * (width - len(q)) for q in seqs]
+            p = self._keep_probs(ids_rows, mask_rows)
+            for (s, e), row in zip(batch, p):
+                k = len(self.pre)
+                probs[s:e] = row[k:k + (e - s)]
+
+        # Score: mean prob over every token touching the atom. Budget: each token is
+        # charged once, to the first atom it touches.
+        sums = [0.0] * len(atoms)
+        hits = [0] * len(atoms)
+        counts = [0] * len(atoms)
+        for a, sets, p in zip(tok_atom, token_atom_sets(offs, owner), probs):
+            if a >= 0:
+                counts[a] += 1
+            for b in sets:
+                sums[b] += p
+                hits[b] += 1
+        scores = [sums[i] / hits[i] if hits[i] else 0.0 for i in range(len(atoms))]
+        return atoms, scores, counts
+
+    def _keep_probs(self, ids_rows: list[list[int]], mask_rows: list[list[int]]) -> list[list[float]]:
+        if self.backend == "onnx":
+            import numpy as np
+
+            logits = self.session.run(["logits"], {"input_ids": np.array(ids_rows, dtype=np.int64),
+                                                   "attention_mask": np.array(mask_rows, dtype=np.int64)})[0]
+            z = logits - logits.max(-1, keepdims=True)
+            e = np.exp(z)
+            return (e[..., 1] / e.sum(-1)).tolist()
+        t = self.torch
+        input_ids = t.tensor(ids_rows, device=self.device)
+        mask = t.tensor(mask_rows, device=self.device)
+        with t.no_grad():
+            logits = self.model(input_ids=input_ids, attention_mask=mask).logits.float()
+        return logits.softmax(-1)[..., 1].cpu().tolist()
+
+    def compress(self, text: str, rate: float | None = 0.5, threshold: float | None = None,
+                 force_protected: bool = False, keep_structure: bool = False, unit="chars",
+                 measure=None) -> dict:
+        """Keep the highest-scoring atoms until `rate` of the text is kept.
+
+        unit="chars" measures the budget in characters (one per whitespace gap), which
+        tracks a downstream LLM's token count without depending on our tokenizer: the
+        encoder splits digits one by one, so in "tokens" identifiers look expensive.
+        unit may also be a callable mapping a list of strings to their costs.
+        measure: optional callable str -> int (e.g. the target LLM's token count). The
+        cut is then binary-searched so measure(output) <= rate * measure(text) exactly,
+        which makes `rate` comparable across compressors with different tokenizers.
+        """
+        atoms, scores, tok_counts = self.atom_scores(text)
+        if callable(unit):
+            counts = list(unit([(" " if a.ws else "") + a.text for a in atoms]))
+        elif unit == "chars":
+            counts = [len(a.text) + (1 if a.ws else 0) for a in atoms]
+        elif unit == "tokens":
+            counts = tok_counts
+        else:
+            raise ValueError(f"unit must be 'chars' or 'tokens', not {unit!r}")
+        forced = [
+            bool(a.text) and ((force_protected and is_protected(a.text)) or (keep_structure and a.text in STRUCTURE))
+            for a in atoms
+        ]
+        if threshold is not None:
+            keep = [f or s >= threshold for f, s in zip(forced, scores)]
+        elif measure is not None and rate is not None:
+            order = [i for i in sorted(range(len(atoms)), key=lambda i: -scores[i])
+                     if atoms[i].text and not forced[i]]
+            target = rate * measure(text)
+
+            def with_top(k: int) -> list[bool]:
+                kk = list(forced)
+                for i in order[:k]:
+                    kk[i] = True
+                return kk
+
+            lo, hi = 0, len(order)  # largest k whose output fits the target
+            while lo < hi:
+                mid = (lo + hi + 1) // 2
+                if measure(detok(atoms, with_top(mid))) <= target:
+                    lo = mid
+                else:
+                    hi = mid - 1
+            keep = with_top(lo)
+        else:
+            total = sum(counts)
+            budget = math.ceil((rate if rate is not None else 1.0) * total)
+            keep = list(forced)
+            used = sum(c for c, f in zip(counts, forced) if f)
+            for i in sorted(range(len(atoms)), key=lambda i: -scores[i]):
+                if used >= budget:
+                    break
+                if not keep[i] and atoms[i].text:
+                    keep[i] = True
+                    used += counts[i]
+        out = detok(atoms, keep)
+        return {"text": out, "kept": sum(c for c, k in zip(counts, keep) if k), "total": sum(counts)}

ctxprune/idcheck.py

@@ -0,0 +1,59 @@
+"""Identifier-survival metric: did the IDs, numbers and paths make it through?
+
+A protected atom survives only if it appears intact as an atom of the
+compressed text, so "0.21" -> "0. 21" counts as lost: that is what the
+downstream LLM actually sees.
+"""
+
+from __future__ import annotations
+
+import json
+
+from .atoms import atomize, is_protected
+
+
+def protected_atoms(text: str) -> set[str]:
+    return {a.text for a in atomize(text) if is_protected(a.text)}
+
+
+def survival(original: str, compressed: str, count_tokens=None) -> dict:
+    want = protected_atoms(original)
+    have = {a.text for a in atomize(compressed)}
+    lost = sorted(want - have)
+    r = {
+        "protected": len(want),
+        "lost": len(lost),
+        "lost_examples": lost[:10],
+        "char_ratio": len(compressed) / max(1, len(original)),
+    }
+    if count_tokens is not None:
+        r["llm_token_ratio"] = count_tokens(compressed) / max(1, count_tokens(original))
+    try:
+        json.loads(original)
+    except (json.JSONDecodeError, ValueError):
+        r["json"] = None
+    else:
+        try:
+            json.loads(compressed)
+            r["json"] = True
+        except (json.JSONDecodeError, ValueError):
+            r["json"] = False
+    return r
+
+
+def summarize(rows: list[dict]) -> dict:
+    prot = sum(r["protected"] for r in rows)
+    lost = sum(r["lost"] for r in rows)
+    js = [r["json"] for r in rows if r.get("json") is not None]
+    summary = {
+        "n": len(rows),
+        "protected": prot,
+        "lost": lost,
+        "lost_rate": lost / prot if prot else 0.0,
+        "json_inputs": len(js),
+        "json_still_valid": sum(js),
+        "mean_char_ratio": sum(r["char_ratio"] for r in rows) / len(rows) if rows else 0.0,
+    }
+    if rows and "llm_token_ratio" in rows[0]:
+        summary["mean_llm_token_ratio"] = sum(r["llm_token_ratio"] for r in rows) / len(rows)
+    return summary

ctxprune/integrations.py

@@ -0,0 +1,90 @@
+"""Drop-in compressors for LangChain and LlamaIndex retrieval pipelines.
+
+    from ctxprune.integrations import CtxpruneDocumentCompressor   # LangChain
+    retriever = ContextualCompressionRetriever(base_compressor=CtxpruneDocumentCompressor(rate=0.5),
+                                               base_retriever=base)
+
+    from ctxprune.integrations import CtxpruneNodePostprocessor    # LlamaIndex
+    engine = index.as_query_engine(node_postprocessors=[CtxpruneNodePostprocessor(rate=0.5)])
+
+Both compress each document's text in place and leave metadata alone. Classes are built
+on first access so neither framework is required to import ctxprune.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+
+from .cli import DEFAULT_MODEL
+
+
+@lru_cache(maxsize=4)
+def _compressor(model: str, backend: str | None):
+    from .compress import Compressor
+
+    if backend is None:
+        try:
+            import onnxruntime  # noqa: F401
+            backend = "onnx"
+        except ImportError:
+            backend = "torch"
+    return Compressor(model, backend=backend)
+
+
+def _langchain():
+    from langchain_core.documents import Document
+    from langchain_core.documents.compressor import BaseDocumentCompressor
+
+    class CtxpruneDocumentCompressor(BaseDocumentCompressor):
+        """LangChain document compressor: keeps `rate` of each document's text."""
+
+        model: str = DEFAULT_MODEL
+        backend: str | None = None
+        rate: float = 0.5
+        force_protected: bool = False
+
+        def compress_documents(self, documents, query, callbacks=None):
+            c = _compressor(self.model, self.backend)
+            return [Document(page_content=c.compress(d.page_content, rate=self.rate,
+                                                     force_protected=self.force_protected)["text"],
+                             metadata=d.metadata, id=d.id)
+                    for d in documents]
+
+    return CtxpruneDocumentCompressor
+
+
+def _llamaindex():
+    from llama_index.core.postprocessor.types import BaseNodePostprocessor
+
+    class CtxpruneNodePostprocessor(BaseNodePostprocessor):
+        """LlamaIndex node postprocessor: keeps `rate` of each retrieved node's text."""
+
+        model: str = DEFAULT_MODEL
+        backend: str | None = None
+        rate: float = 0.5
+        force_protected: bool = False
+
+        @classmethod
+        def class_name(cls) -> str:
+            return "CtxpruneNodePostprocessor"
+
+        def _postprocess_nodes(self, nodes, query_bundle=None):
+            c = _compressor(self.model, self.backend)
+            for n in nodes:
+                n.node.set_content(c.compress(n.node.get_content(), rate=self.rate,
+                                              force_protected=self.force_protected)["text"])
+            return nodes
+
+    return CtxpruneNodePostprocessor
+
+
+_FACTORIES = {"CtxpruneDocumentCompressor": _langchain, "CtxpruneNodePostprocessor": _llamaindex}
+_built: dict[str, type] = {}
+
+
+def __getattr__(name: str):
+    if name not in _FACTORIES:
+        raise AttributeError(name)
+    if name not in _built:
+        _built[name] = _FACTORIES[name]()
+    return _built[name]