skhavin/proactive-cache

1"""2quickstart.py — Minimal 10-line example of ProactiveCache.3 4Shows how to apply O(n) generation to any HuggingFace model.5"""6 7from transformers import AutoModelForCausalLM, AutoTokenizer8from proactive_cache import ProactiveCache9 10MODEL = "meta-llama/Llama-3.1-8B"   # replace with any HF model11 12# Load model (any HuggingFace CausalLM)13tokenizer = AutoTokenizer.from_pretrained(MODEL)14model = AutoModelForCausalLM.from_pretrained(MODEL, device_map="auto")15 16# ── Step 1: Apply O(n) eviction (one line) ───────────────────────────────────17model = ProactiveCache.apply(model, budget=512)18 19# ── Step 2: Profile once on calibration data (saves proactive_cache_prototypes.pkl)20ProactiveCache.profile(model, tokenizer, corpus="wikitext", num_docs=50)21 22# ── Step 3: All inference is now O(n) ────────────────────────────────────────23prompt = "In the age of long-context language models,"24inputs = tokenizer(prompt, return_tensors="pt").to(model.device)25output = model.generate(**inputs, max_new_tokens=200, do_sample=False)26print(tokenizer.decode(output[0], skip_special_tokens=True))27