kami
047e2a4070
- build() is now suspend: pipeline runs all stages in fixed order, each gated by level - TOKEN_PRUNE (level 3): TokenPruner interface + LLMLingua-2 sidecar (sidecars/llmlingua) + HttpTokenPruner adapter (fails open if sidecar down); prunes freeform, preserves protected spans, skips tier-0 turns when TIER_SPLIT on - TOME_MERGE (level 8): ToMeMerger collapses near-duplicate freeform turns (Jaccard) - Stage 5 selection: RelevanceScorer + EmbeddingRelevanceScorer (cosine over Embedder); query-conditioned reorder so least-relevant freeform drops first under budget - [orchestration] compression_level + token_pruner_url config, wired in Main - suspend ripple fixed across builder callers/stubs
66 lines
1.9 KiB
Python
66 lines
1.9 KiB
Python
"""LLMLingua-2 token-pruning sidecar for correx.
|
|
|
|
A tiny HTTP service correx calls (via HttpTokenPruner) to prune low-perplexity tokens from
|
|
freeform prose before it's sent to the local LLM. Kept in Python because LLMLingua-2 is a
|
|
torch/BERT classifier with no JVM equivalent (pipeline §4).
|
|
|
|
Contract:
|
|
POST /prune {"text": str, "protected": [str], "rate": float} -> {"compressed": str}
|
|
rate = fraction of tokens to KEEP (0.55 keeps ~55%, i.e. ~45% compression).
|
|
`protected` substrings are force-kept verbatim (IDs, numbers, paths, code).
|
|
GET /health -> {"status": "ok"}
|
|
|
|
Run:
|
|
pip install -r requirements.txt
|
|
uvicorn server:app --host 127.0.0.1 --port 8199
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from fastapi import FastAPI
|
|
from pydantic import BaseModel
|
|
|
|
app = FastAPI(title="correx-llmlingua")
|
|
|
|
_MODEL = os.environ.get("LLMLINGUA_MODEL", "microsoft/llmlingua-2-xlm-roberta-large-meetingbank")
|
|
_compressor = None
|
|
|
|
|
|
def _get_compressor():
|
|
# Lazy-load so /health works (and the process starts fast) before torch spins up.
|
|
global _compressor
|
|
if _compressor is None:
|
|
from llmlingua import PromptCompressor
|
|
_compressor = PromptCompressor(model_name=_MODEL, use_llmlingua2=True)
|
|
return _compressor
|
|
|
|
|
|
class PruneRequest(BaseModel):
|
|
text: str
|
|
protected: list[str] = []
|
|
rate: float = 0.55
|
|
|
|
|
|
class PruneResponse(BaseModel):
|
|
compressed: str
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.post("/prune", response_model=PruneResponse)
|
|
def prune(req: PruneRequest):
|
|
text = req.text.strip()
|
|
if not text:
|
|
return PruneResponse(compressed=req.text)
|
|
# rate is fraction to keep; LLMLingua-2 force_tokens keeps the protected spans verbatim.
|
|
result = _get_compressor().compress_prompt(
|
|
text,
|
|
rate=max(0.1, min(1.0, req.rate)),
|
|
force_tokens=req.protected or None,
|
|
drop_consecutive=True,
|
|
)
|
|
return PruneResponse(compressed=result["compressed_prompt"])
|