prompt-compress 0.1.0__tar.gz

prompt_compress-0.1.0/.gitignore

@@ -0,0 +1,30 @@
+# Secrets
+.env
+
+# Python bytecode / packaging
+__pycache__/
+*.py[cod]
+*$py.class
+*.egg-info/
+*.egg
+dist/
+build/
+.eggs/
+
+# Virtualenvs
+.venv/
+venv/
+env/
+
+# OS / editor cruft
+.DS_Store
+.idea/
+.vscode/
+
+# Frontend dependencies / build metadata
+node_modules/
+*.tsbuildinfo
+
+# Research caches (regenerate via research/benchmark.py + research/evaluate.py)
+data/results/benchmark/*.json
+!data/results/benchmark/.gitkeep

prompt_compress-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 joela03
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

prompt_compress-0.1.0/PKG-INFO

@@ -0,0 +1,133 @@
+Metadata-Version: 2.4
+Name: prompt-compress
+Version: 0.1.0
+Summary: Structural prompt compression with safety gating
+Project-URL: Homepage, https://github.com/joela03/bayesian-prompt-compressor-
+Project-URL: Repository, https://github.com/joela03/bayesian-prompt-compressor-
+Project-URL: Issues, https://github.com/joela03/bayesian-prompt-compressor-/issues
+Author: joela03
+License: MIT
+License-File: LICENSE
+Keywords: bayesian-optimization,compression,llm,prompt
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Requires-Dist: networkx
+Requires-Dist: numpy
+Requires-Dist: scikit-learn
+Requires-Dist: sentence-transformers
+Provides-Extra: openai
+Requires-Dist: openai>=1.0; extra == 'openai'
+Provides-Extra: research
+Requires-Dist: datasets; extra == 'research'
+Requires-Dist: llmlingua>=0.2; extra == 'research'
+Requires-Dist: matplotlib; extra == 'research'
+Requires-Dist: openai>=1.0; extra == 'research'
+Requires-Dist: pandas; extra == 'research'
+Requires-Dist: python-dotenv; extra == 'research'
+Requires-Dist: scipy; extra == 'research'
+Requires-Dist: seaborn; extra == 'research'
+Description-Content-Type: text/markdown
+
+# prompt-compress
+
+Structural prompt compression for production LLM apps. Where LLMLingua removes individual low-perplexity tokens, this library parses your system prompt into named components (instruction, examples, constraints, style, context), uses Bayesian optimisation to search which components to keep and how aggressively to compress each, scores candidates by semantic similarity to the original, and gates every output through a post-compression validator (persona / placeholder / similarity). Prompts that are already information-dense are detected up front and passed through unchanged.
+
+## Install
+
+```bash
+pip install prompt-compress
+```
+
+## Quickstart — production integration
+
+```python
+from prompt_compress import PromptCompressor, CompressionFailedError
+
+compressor = PromptCompressor()
+
+try:
+    result = compressor.compress(
+        SYSTEM_PROMPT,
+        min_similarity=0.80,
+        on_failure='raise',
+    )
+    SYSTEM_PROMPT = result.compressed_text
+    print(f"Saved {result.tokens_saved} tokens per call ({result.compression_ratio:.1%})")
+except CompressionFailedError as e:
+    print(f"Compression unsafe, using original: {e}")
+```
+
+`on_failure` accepts `'fallback'` (default — return the original silently with `gate_passed=False`), `'raise'` (raise `CompressionFailedError`), or `'warn'` (log a warning and return the fallback). The library never blocks on user input.
+
+## Inspecting results
+
+```python
+result = compressor.compress(SYSTEM_PROMPT)
+
+print(result.summary())   # one-screen terminal summary
+print(result.diff())      # side-by-side original vs compressed
+result.to_dict()          # JSON-serialisable, useful for caching/logging
+```
+
+Key properties on `CompressionResult`:
+
+| Property | Description |
+|---|---|
+| `compressed_text` | the output you should use |
+| `compression_ratio` | tokens saved / original tokens |
+| `tokens_saved` | absolute token count saved |
+| `semantic_similarity` | cosine sim of original vs compressed (MiniLM) |
+| `compression_efficiency` | `compression_ratio × semantic_similarity` |
+| `safe_to_use` | True iff all validator checks passed |
+| `persona_preserved` | True iff the "You are…" line survived |
+| `placeholders_preserved` | True iff every `{var}` from the original is in the output |
+| `tier` / `tier_label` | which pipeline tier ran (1 BO, 2 TextRank, 3 Preserved) |
+| `density` | information density score used for routing |
+
+## Configuration
+
+```python
+from prompt_compress import PromptCompressor, OptimisationConfig
+
+compressor = PromptCompressor(
+    # Optimiser variants:
+    use_informed_prior=False,    # seed BO with P3-derived prior
+    use_attention_prior=False,   # per-prompt attention prior + ISR safety gate
+    # Trade-off knob:
+    alpha=0.3,                   # "auto" → 0.3 (validated benchmark default)
+    # Tune BO budget:
+    optimisation_config=OptimisationConfig(
+        n_iterations=20, n_init=5, beta=2.0, random_seed=42,
+    ),
+)
+```
+
+`min_similarity` and `on_failure` are per-call (`compressor.compress(prompt, min_similarity=…, on_failure=…)`) so different parts of your app can adopt different safety bars without rebuilding the compressor.
+
+## Benchmark results
+
+Matched-subset comparison against LLMLingua on the 38 prompts both systems successfully compressed (see `research/benchmark.py` and `research/evaluate.py` to reproduce):
+
+| Metric                       | Ours    | LLMLingua |
+|------------------------------|---------|-----------|
+| Compression ratio            | 24.1%   | 24.2%     |
+| LLM judge score (0–100)      | 73.3    | 70.2      |
+| Persona preservation         | 100%    | 53%       |
+| Compression efficiency       | 0.179   | 0.155     |
+
+`Compression efficiency = compression_ratio × output_similarity` — rewards being high on both axes.
+
+## Citation
+
+> *EMNLP manuscript in preparation.*

prompt_compress-0.1.0/README.md

@@ -0,0 +1,92 @@
+# prompt-compress
+
+Structural prompt compression for production LLM apps. Where LLMLingua removes individual low-perplexity tokens, this library parses your system prompt into named components (instruction, examples, constraints, style, context), uses Bayesian optimisation to search which components to keep and how aggressively to compress each, scores candidates by semantic similarity to the original, and gates every output through a post-compression validator (persona / placeholder / similarity). Prompts that are already information-dense are detected up front and passed through unchanged.
+
+## Install
+
+```bash
+pip install prompt-compress
+```
+
+## Quickstart — production integration
+
+```python
+from prompt_compress import PromptCompressor, CompressionFailedError
+
+compressor = PromptCompressor()
+
+try:
+    result = compressor.compress(
+        SYSTEM_PROMPT,
+        min_similarity=0.80,
+        on_failure='raise',
+    )
+    SYSTEM_PROMPT = result.compressed_text
+    print(f"Saved {result.tokens_saved} tokens per call ({result.compression_ratio:.1%})")
+except CompressionFailedError as e:
+    print(f"Compression unsafe, using original: {e}")
+```
+
+`on_failure` accepts `'fallback'` (default — return the original silently with `gate_passed=False`), `'raise'` (raise `CompressionFailedError`), or `'warn'` (log a warning and return the fallback). The library never blocks on user input.
+
+## Inspecting results
+
+```python
+result = compressor.compress(SYSTEM_PROMPT)
+
+print(result.summary())   # one-screen terminal summary
+print(result.diff())      # side-by-side original vs compressed
+result.to_dict()          # JSON-serialisable, useful for caching/logging
+```
+
+Key properties on `CompressionResult`:
+
+| Property | Description |
+|---|---|
+| `compressed_text` | the output you should use |
+| `compression_ratio` | tokens saved / original tokens |
+| `tokens_saved` | absolute token count saved |
+| `semantic_similarity` | cosine sim of original vs compressed (MiniLM) |
+| `compression_efficiency` | `compression_ratio × semantic_similarity` |
+| `safe_to_use` | True iff all validator checks passed |
+| `persona_preserved` | True iff the "You are…" line survived |
+| `placeholders_preserved` | True iff every `{var}` from the original is in the output |
+| `tier` / `tier_label` | which pipeline tier ran (1 BO, 2 TextRank, 3 Preserved) |
+| `density` | information density score used for routing |
+
+## Configuration
+
+```python
+from prompt_compress import PromptCompressor, OptimisationConfig
+
+compressor = PromptCompressor(
+    # Optimiser variants:
+    use_informed_prior=False,    # seed BO with P3-derived prior
+    use_attention_prior=False,   # per-prompt attention prior + ISR safety gate
+    # Trade-off knob:
+    alpha=0.3,                   # "auto" → 0.3 (validated benchmark default)
+    # Tune BO budget:
+    optimisation_config=OptimisationConfig(
+        n_iterations=20, n_init=5, beta=2.0, random_seed=42,
+    ),
+)
+```
+
+`min_similarity` and `on_failure` are per-call (`compressor.compress(prompt, min_similarity=…, on_failure=…)`) so different parts of your app can adopt different safety bars without rebuilding the compressor.
+
+## Benchmark results
+
+Matched-subset comparison against LLMLingua on the 38 prompts both systems successfully compressed (see `research/benchmark.py` and `research/evaluate.py` to reproduce):
+
+| Metric                       | Ours    | LLMLingua |
+|------------------------------|---------|-----------|
+| Compression ratio            | 24.1%   | 24.2%     |
+| LLM judge score (0–100)      | 73.3    | 70.2      |
+| Persona preservation         | 100%    | 53%       |
+| Compression efficiency       | 0.179   | 0.155     |
+
+`Compression efficiency = compression_ratio × output_similarity` — rewards being high on both axes.
+
+## Citation
+
+> *EMNLP manuscript in preparation.*

prompt_compress-0.1.0/prompt_compress/__init__.py

@@ -0,0 +1,26 @@
+"""
+prompt-compress — structural prompt compression with safety gating.
+
+Public API:
+    PromptCompressor       — the compressor
+    CompressionResult      — what compress() returns
+    CompressionFailedError — raised when on_failure='raise' and validation fails
+    OptimisationConfig     — knobs for the Bayesian optimiser
+
+Internal classes (parser, validator, evaluator, etc.) are importable from
+their submodules for advanced use but are intentionally not promoted here.
+"""
+
+from .compressor import CompressionFailedError, PromptCompressor
+from .optimiser import OptimisationConfig
+from .result import CompressionResult
+
+__version__ = '0.1.0'
+
+__all__ = [
+    'PromptCompressor',
+    'CompressionFailedError',
+    'OptimisationConfig',
+    'CompressionResult',
+    '__version__',
+]

prompt_compress-0.1.0/prompt_compress/_persona.py

@@ -0,0 +1,27 @@
+"""
+Shared persona-detection patterns.
+
+Lives in its own module so both `parser.py` and `validators.py` can import
+from it without introducing a circular dependency between them.
+"""
+
+import re
+
+PERSONA_PATTERNS = (
+    'you are',
+    'act as',
+    'i want you to act',
+    'pretend to be',
+    'pretend you are',
+    'imagine you are',
+    "let's role[- ]?play as",
+    'you will (be|act as|play)',
+    'assume the role',
+)
+
+
+def persona_present(text: str) -> bool:
+    """True if text opens with a recognised persona pattern."""
+    first_line = text.lstrip().split('\n', 1)[0].strip()
+    first_sentence = re.split(r'[.!?]', first_line, maxsplit=1)[0].strip().lower()
+    return any(re.match(rf'^{pat}\b', first_sentence) for pat in PERSONA_PATTERNS)

prompt_compress-0.1.0/prompt_compress/attention_optimiser.py

@@ -0,0 +1,201 @@
+"""
+Attention-informed Bayesian optimiser.
+
+Subclass of InformedBayesianOptimiser that swaps the static P3-JSON prior
+for a per-prompt prior generated from component-to-component attention,
+and adds an Information Sufficiency Ratio (ISR) pre-compression gate so
+prompts already near their Minimum Description Length are not compressed.
+"""
+
+import logging
+from typing import Dict, Optional
+
+from .attention_priors import AttentionPriorGenerator
+from .encoders import PromptStructure
+from .information_sufficiency import ISRGate
+from .informed_optimiser import InformedBayesianOptimiser
+from .optimiser import OptimisationConfig, OptimisationResult
+
+logger = logging.getLogger(__name__)
+
+
+class AttentionInformedOptimiser(InformedBayesianOptimiser):
+    """
+    Use a fresh attention-derived prior for each prompt instead of a static
+    dataset-wide prior, with an ISR gate that can short-circuit BO when the
+    prompt is information-dense (near MDL) or flag aggressive compression
+    when the prompt is highly redundant.
+    """
+
+    AGGRESSIVE_BETA_MULTIPLIER = 1.2  # widen UCB exploration when ISR is low
+    AUTO_ALPHA = 0.3                  # default for alpha="auto"; matches the
+                                      # validated benchmark (see SemanticEvaluator.AUTO_ALPHA)
+
+    def __init__(
+        self,
+        encoder,
+        evaluator,
+        config: Optional[OptimisationConfig] = None,
+        components: Optional[Dict] = None,
+        prior_generator: Optional[AttentionPriorGenerator] = None,
+        prompt_text: Optional[str] = None,
+        use_isr_gate: bool = True,
+        isr_high_threshold: float = 0.85,
+        isr_low_threshold: float = 0.40,
+        alpha=AUTO_ALPHA,
+    ):
+        """
+        Args:
+            components:         Parsed components dict (PromptParser output) for
+                                the current prompt. Required to generate the prior.
+            prior_generator:    Optional shared AttentionPriorGenerator instance.
+            prompt_text:        Raw prompt text. Required for the ISR gate to
+                                fire; if None, the gate is silently skipped.
+            use_isr_gate:       Master switch for the ISR pre-check.
+            isr_high_threshold: ISR above this ⇒ skip compression entirely.
+            isr_low_threshold:  ISR below this ⇒ enable aggressive mode (the
+                                UCB beta is multiplied by AGGRESSIVE_BETA_MULTIPLIER
+                                to favour exploration of compressive structures).
+            alpha:              Quality/compression trade-off. Larger alpha
+                                penalises length more heavily. Accepts "auto"
+                                (resolves to AUTO_ALPHA=1.0) or any float.
+                                The objective lives on SemanticEvaluator —
+                                this kwarg is stored for logging and is
+                                surfaced on the OptimisationResult.
+        """
+        self.alpha = self._resolve_alpha(alpha)
+        if components is None:
+            raise ValueError(
+                "AttentionInformedOptimiser requires the parsed components dict"
+            )
+
+        self._components = components
+        self._prior_generator = prior_generator or AttentionPriorGenerator()
+        self._prompt_text = prompt_text
+
+        if config is None:
+            config = OptimisationConfig(
+                n_iterations=20, n_init=5, beta=2.0, random_seed=42
+            )
+
+        # Honour both the explicit kwargs and the OptimisationConfig fields.
+        # Explicit kwargs win so callers can override per-instance.
+        self.use_isr_gate = use_isr_gate and config.enable_isr_gate
+        self.isr_gate = ISRGate(
+            high_threshold=isr_high_threshold if isr_high_threshold is not None else config.isr_high,
+            low_threshold=isr_low_threshold if isr_low_threshold is not None else config.isr_low,
+        )
+
+        # Skip the parent's JSON-loading constructor by initialising the
+        # grandparent (BayesianPromptOptimiser) directly.
+        from .optimiser import BayesianPromptOptimiser
+        BayesianPromptOptimiser.__init__(self, encoder, evaluator, config)
+
+        self.prior = self._prior_generator.generate(components)
+        logger.info(
+            "Loaded attention-informed prior "
+            "(mean_attention=%s)",
+            {k: round(v, 2) for k, v in self.prior['mean_attention'].items()},
+        )
+
+    @classmethod
+    def _resolve_alpha(cls, alpha) -> float:
+        if isinstance(alpha, str):
+            if alpha == "auto":
+                return cls.AUTO_ALPHA
+            raise ValueError(f"alpha must be 'auto' or numeric, got {alpha!r}")
+        return float(alpha)
+
+    def check_isr(self) -> tuple[bool, float, str]:
+        """
+        Run the ISR gate against the stored prompt text. Returns the same
+        triple as ISRGate.should_compress.
+
+        Returns (True, 0.0, "isr disabled") if the gate is off or no prompt
+        text was supplied.
+        """
+        if not self.use_isr_gate or self._prompt_text is None:
+            return True, 0.0, "isr disabled"
+        return self.isr_gate.should_compress(self._prompt_text)
+
+    def optimise(self) -> OptimisationResult:
+        """
+        Run ISR gate first, then dispatch to the parent BO loop.
+
+        Behaviour:
+          - ISR > high_threshold: return a `skipped=True` result; the
+            orchestrator (PromptCompressor) treats this as "return original".
+          - ISR < low_threshold:  temporarily inflate UCB beta to favour
+            exploration of compressive structures.
+          - Otherwise:           run BO unchanged.
+        """
+        should_compress, isr, reason = self.check_isr()
+
+        if not should_compress:
+            logger.info(
+                "ISR=%.2f exceeds threshold. Prompt already near Minimum "
+                "Description Length. Preserving original.",
+                isr,
+            )
+            return self._skipped_result(isr, reason)
+
+        if self.use_isr_gate and self._prompt_text is not None:
+            band = "low" if isr < self.isr_gate.low_threshold else (
+                "high" if isr > self.isr_gate.high_threshold else "moderate"
+            )
+            logger.info("ISR=%.2f (%s). %s.", isr, band, reason.capitalize())
+
+        # Aggressive mode: temporarily widen UCB exploration. Restore beta
+        # after the run so the optimiser stays reusable.
+        original_beta = self.config.beta
+        if isr < self.isr_gate.low_threshold:
+            self.config.beta = original_beta * self.AGGRESSIVE_BETA_MULTIPLIER
+            logger.info(
+                "ISR=%.2f (low). Enabling aggressive compression "
+                "(beta %.2f -> %.2f).",
+                isr,
+                original_beta,
+                self.config.beta,
+            )
+
+        try:
+            result = super().optimise()
+        finally:
+            self.config.beta = original_beta
+
+        result.skipped = False
+        result.isr_score = isr if self.use_isr_gate and self._prompt_text else None
+        result.isr_reason = reason if self.use_isr_gate and self._prompt_text else None
+        result.alpha_used = self.alpha
+        return result
+
+    def _skipped_result(self, isr: float, reason: str) -> OptimisationResult:
+        """
+        Build a sentinel OptimisationResult signalling "do not compress".
+
+        The structure is set to "keep everything" so that if a caller ignores
+        `skipped` and materialises this structure anyway, the output is at
+        worst the full original prompt — not a corrupted compression.
+        """
+        identity_structure = PromptStructure(
+            has_instruction=True,
+            has_examples=True,
+            has_constraints=True,
+            has_style=True,
+            has_context=True,
+            num_examples=1.0,
+            instruction_length=1.0,
+            total_tokens=1.0,
+            component_ordering=[1, 2, 3, 4, 5],
+        )
+        return OptimisationResult(
+            best_structure=identity_structure,
+            best_score=1.0,
+            all_scores=[1.0],
+            all_structures=[identity_structure],
+            total_evaluations=0,
+            skipped=True,
+            isr_score=isr,
+            isr_reason=reason,
+            alpha_used=self.alpha,
+        )

prompt_compress-0.1.0/prompt_compress/attention_priors.py

@@ -0,0 +1,149 @@
+"""
+Per-prompt attention-informed priors for the Bayesian optimiser.
+
+For each component in the parsed prompt, compute the mean cosine similarity
+("attention") to every other component using sentence-transformers embeddings.
+High mean attention (>0.7) means the component is tightly coupled to the rest
+of the prompt — keep it. Low mean attention (<0.4) means it is independent —
+safe to drop.
+
+The output is shaped to match `InformedBayesianOptimiser._load_prior()` so the
+existing prior pipeline can consume it without changes.
+"""
+
+import numpy as np
+from typing import Dict, Optional
+from sentence_transformers import SentenceTransformer
+
+COMPONENT_NAMES = ['instruction', 'examples', 'constraints', 'style', 'context']
+
+
+class AttentionPriorGenerator:
+    """
+    Generate a prompt-specific prior from component-to-component attention.
+    """
+
+    EMBEDDING_MODEL = 'all-MiniLM-L6-v2'
+
+    HIGH_THRESHOLD = 0.7
+    LOW_THRESHOLD = 0.4
+
+    # Inclusion probabilities mapped from coupling strength
+    HIGH_INCLUSION = 0.85
+    MID_INCLUSION = 0.5
+    LOW_INCLUSION = 0.25
+
+    _shared_model: Optional[SentenceTransformer] = None
+
+    def __init__(self):
+        if AttentionPriorGenerator._shared_model is None:
+            AttentionPriorGenerator._shared_model = SentenceTransformer(self.EMBEDDING_MODEL)
+        self.model = AttentionPriorGenerator._shared_model
+
+    def generate(self, components: Dict) -> Dict:
+        """
+        Build a prior dict compatible with InformedBayesianOptimiser.
+
+        Args:
+            components: Dict[str, list[str]] from PromptParser.parse(). May
+                        include a '__protected__' key, which is ignored.
+
+        Returns:
+            Dict with keys 'mean', 'variance', 'source', plus an extra
+            'attention_matrix' (5x5) and 'mean_attention' (per-component)
+            for debugging/inspection.
+        """
+        present, embeddings = self._embed_components(components)
+
+        # Default inclusion probabilities (used when a component is empty)
+        prior_mean = {
+            'has_instruction': self.HIGH_INCLUSION,  # always favour instruction
+            'has_examples': self.MID_INCLUSION,
+            'has_constraints': self.MID_INCLUSION,
+            'has_style': self.LOW_INCLUSION,
+            'has_context': self.LOW_INCLUSION,
+            'instruction_length': 0.7,
+            'num_examples': 0.4,
+        }
+
+        attention_matrix = np.zeros((5, 5))
+        mean_attention = {name: 0.0 for name in COMPONENT_NAMES}
+
+        if len(present) >= 2:
+            sim = self._cosine_matrix(embeddings)
+            for i, name_i in enumerate(present):
+                col = COMPONENT_NAMES.index(name_i)
+                for j, name_j in enumerate(present):
+                    row = COMPONENT_NAMES.index(name_j)
+                    attention_matrix[col, row] = sim[i, j]
+                # Mean attention excludes the diagonal (self-similarity = 1.0)
+                others = [sim[i, j] for j in range(len(present)) if j != i]
+                mean_attention[name_i] = float(np.mean(others)) if others else 0.0
+
+            for name in present:
+                ma = mean_attention[name]
+                if ma >= self.HIGH_THRESHOLD:
+                    prior_mean[f'has_{name}'] = self.HIGH_INCLUSION
+                elif ma < self.LOW_THRESHOLD:
+                    prior_mean[f'has_{name}'] = self.LOW_INCLUSION
+                else:
+                    prior_mean[f'has_{name}'] = self.MID_INCLUSION
+
+        # Variance: shrink when we have strong signal (many present components),
+        # widen when we have little to learn from.
+        prior_variance = 0.15 if len(present) >= 3 else 0.25
+
+        return {
+            'mean': prior_mean,
+            'variance': prior_variance,
+            'source': 'per-prompt attention',
+            'attention_matrix': attention_matrix.tolist(),
+            'mean_attention': mean_attention,
+        }
+
+    def _embed_components(self, components: Dict):
+        """
+        Embed each present component (concatenated sentences) and return the
+        ordered list of names and their embedding matrix.
+        """
+        present_names: list[str] = []
+        texts: list[str] = []
+        for name in COMPONENT_NAMES:
+            sentences = components.get(name, [])
+            if not sentences:
+                continue
+            joined = ' '.join(sentences) if isinstance(sentences, list) else str(sentences)
+            if not joined.strip():
+                continue
+            present_names.append(name)
+            texts.append(joined)
+        if not texts:
+            return [], np.zeros((0, 384))
+        embeddings = self.model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
+        return present_names, embeddings
+
+    @staticmethod
+    def _cosine_matrix(embeddings: np.ndarray) -> np.ndarray:
+        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
+        norms = np.where(norms == 0.0, 1.0, norms)
+        normalised = embeddings / norms
+        return normalised @ normalised.T
+
+
+if __name__ == "__main__":
+    sample_components = {
+        'instruction': ['You are a SQL tutor. Explain query plans.'],
+        'examples': ['Example: SELECT * FROM users WHERE id = 1.'],
+        'constraints': ['Always include the cost estimate.', 'Never invent table names.'],
+        'style': ['Use a formal tone.'],
+        'context': [],
+    }
+    gen = AttentionPriorGenerator()
+    prior = gen.generate(sample_components)
+    print('mean inclusion probabilities:')
+    for k, v in prior['mean'].items():
+        print(f'  {k}: {v:.3f}')
+    print()
+    print('mean attention per component:')
+    for k, v in prior['mean_attention'].items():
+        print(f'  {k}: {v:.3f}')