From ff8560ae9a6fea658ac1b8fce29cc31b580fb545 Mon Sep 17 00:00:00 2001 From: Joshua Farrow <41028474+warrofua@users.noreply.github.com> Date: Sat, 25 Oct 2025 11:11:36 -0400 Subject: [PATCH 1/2] Allow FUNSD benchmark to customize budget allocation --- benchmarks/doc_understanding.py | 24 +++++++++--- nd_llm/bottleneck/ib.py | 66 ++++++++++++++++++++++++++++++++- 2 files changed, 84 insertions(+), 6 deletions(-) diff --git a/benchmarks/doc_understanding.py b/benchmarks/doc_understanding.py index fc712de..1fd314c 100644 --- a/benchmarks/doc_understanding.py +++ b/benchmarks/doc_understanding.py @@ -30,7 +30,7 @@ except Exception: # pragma: no cover - fallback when torch is unavailable torch = None # type: ignore[assignment] -from nd_llm.bottleneck import CompressionResult, IBottleneck +from nd_llm.bottleneck import CompressionResult, IBottleneck, RegistryAwareBudgetAllocator from nd_llm.orchestration import Orchestrator, UsageEvent from nd_llm.stm import STM from nd_llm.utils import ( @@ -158,8 +158,16 @@ def run_funsd_benchmark( data_root: Optional[Path | str] = None, use_sample: bool = True, seed: int = 0, + allocator_kwargs: Optional[Mapping[str, Any]] = None, ) -> Dict[str, Any]: - """Evaluate FUNSD documents for numeric-answer retention under budget constraints.""" + """Evaluate FUNSD documents for numeric-answer retention under budget constraints. + + Parameters + ---------- + allocator_kwargs: + Optional mapping forwarded to :class:`RegistryAwareBudgetAllocator` so callers can + specify advanced controls such as ``field_weights`` or ``field_min_quota``. + """ registry = build_funsd_registry() build_funsd_encoders(registry) @@ -183,6 +191,7 @@ def run_funsd_benchmark( seed=seed, ablations=ablations, mi_field_priorities=("text", "layout"), + allocator_kwargs=allocator_kwargs, ) runs.append(budget_run) @@ -254,15 +263,17 @@ def _evaluate_budget( seed: int, ablations: Optional[Mapping[str, "AblationFn"]] = None, mi_field_priorities: Optional[Sequence[str]] = None, + allocator_kwargs: Optional[Mapping[str, Any]] = None, ) -> BudgetRun: - bottleneck = IBottleneck(target_budget=int(budget)) + budget_allocator = RegistryAwareBudgetAllocator(**dict(allocator_kwargs or {})) + bottleneck = IBottleneck(target_budget=int(budget), budget_allocator=budget_allocator) ablation_totals: Dict[str, Dict[str, Any]] = {} ablation_bottlenecks: Dict[str, IBottleneck] = {} if ablations: ablation_totals = {name: _make_ablation_totals() for name in ablations} ablation_bottlenecks = { - name: IBottleneck(target_budget=int(budget)) for name in ablations + name: IBottleneck(target_budget=int(budget), budget_allocator=budget_allocator) } with TemporaryDirectory(prefix="ndllm-bench-") as tmp: @@ -371,7 +382,10 @@ def _evaluate_budget( ) ab_bottleneck = ablation_bottlenecks.get(name) if ab_bottleneck is None: - ab_bottleneck = IBottleneck(target_budget=int(budget)) + ab_bottleneck = IBottleneck( + target_budget=int(budget), + budget_allocator=budget_allocator, + ) ablation_bottlenecks[name] = ab_bottleneck ab_result = ab_bottleneck.compress( diff --git a/nd_llm/bottleneck/ib.py b/nd_llm/bottleneck/ib.py index bfaf868..6667f48 100644 --- a/nd_llm/bottleneck/ib.py +++ b/nd_llm/bottleneck/ib.py @@ -167,6 +167,8 @@ def __init__( salience_bonus: float = 1.5, key_weight: float = 0.3, min_weight: float = 0.05, + field_weights: Optional[Mapping[str, Any]] = None, + field_min_quota: Optional[Mapping[str, Any]] = None, ) -> None: if salience_bonus <= 0: raise ValueError("salience_bonus must be positive") @@ -177,6 +179,24 @@ def __init__( self.salience_bonus = salience_bonus self.key_weight = key_weight self.min_weight = min_weight + self.field_weights: Dict[str, float] = {} + for name, value in (field_weights or {}).items(): + try: + weight = float(value) + except (TypeError, ValueError): + continue + if weight <= 0: + continue + self.field_weights[str(name)] = weight + self.field_min_quota: Dict[str, int] = {} + for name, value in (field_min_quota or {}).items(): + try: + quota = int(value) + except (TypeError, ValueError): + continue + if quota <= 0: + continue + self.field_min_quota[str(name)] = quota def __call__( self, @@ -188,10 +208,21 @@ def __call__( available = sum(token_counts.values()) limit = min(int(total_budget), available) weights: Dict[str, float] = {} + quotas: Dict[str, int] = {} + if self.field_min_quota: + for field in scores: + quota = self.field_min_quota.get(field, 0) + if quota <= 0: + continue + quotas[field] = min(int(quota), token_counts.get(field, 0)) + for field in scores: info = metadata.get(field, {}) budget_weight = info.get("budget_weight") if isinstance(info, Mapping) else None - if isinstance(budget_weight, (int, float)): + override_weight = self.field_weights.get(field) + if override_weight is not None: + weight = float(override_weight) + elif isinstance(budget_weight, (int, float)): weight = float(budget_weight) else: keys = info.get("keys", []) if isinstance(info, Mapping) else [] @@ -255,6 +286,39 @@ def __call__( allocations[field] += 1 assigned += 1 + if quotas: + donors: Dict[str, int] = { + field: allocations.get(field, 0) - quotas.get(field, 0) for field in scores + } + for field, quota in quotas.items(): + current = allocations.get(field, 0) + if quota <= 0 or current >= quota: + continue + needed = quota - current + remaining_capacity = token_counts.get(field, 0) - current + if remaining_capacity <= 0: + continue + needed = min(needed, remaining_capacity) + while needed > 0: + donor_field = None + donor_surplus = 0 + for candidate, surplus in donors.items(): + if candidate == field or surplus <= 0: + continue + if surplus > donor_surplus or ( + surplus == donor_surplus and donor_field is not None and candidate < donor_field + ): + donor_field = candidate + donor_surplus = surplus + if donor_field is None or donor_surplus <= 0: + break + transfer = min(needed, donor_surplus) + allocations[donor_field] -= transfer + donors[donor_field] -= transfer + allocations[field] += transfer + donors[field] = allocations[field] - quotas.get(field, 0) + needed -= transfer + return allocations, {field: weights.get(field, 0.0) for field in scores} From ead18c7188db9d62bf2e228f60d59b4ca2b8f34d Mon Sep 17 00:00:00 2001 From: Joshua Farrow <41028474+warrofua@users.noreply.github.com> Date: Sat, 25 Oct 2025 11:18:41 -0400 Subject: [PATCH 2/2] Fix ablation bottleneck creation comprehension --- benchmarks/doc_understanding.py | 1 + 1 file changed, 1 insertion(+) diff --git a/benchmarks/doc_understanding.py b/benchmarks/doc_understanding.py index 1fd314c..fe8ff1c 100644 --- a/benchmarks/doc_understanding.py +++ b/benchmarks/doc_understanding.py @@ -274,6 +274,7 @@ def _evaluate_budget( ablation_totals = {name: _make_ablation_totals() for name in ablations} ablation_bottlenecks = { name: IBottleneck(target_budget=int(budget), budget_allocator=budget_allocator) + for name in ablations } with TemporaryDirectory(prefix="ndllm-bench-") as tmp: