greengate 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Tharusha Rasath Hemachandra
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,136 @@
1
+ Metadata-Version: 2.4
2
+ Name: greengate
3
+ Version: 0.1.0
4
+ Summary: Confidence-aware cascading for green AI inference: route easy queries to a small local model, escalate only when uncertain, and account for the carbon honestly
5
+ Author: Tharusha Rasath Hemachandra
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/Rasath16/GreenGate
8
+ Project-URL: Repository, https://github.com/Rasath16/GreenGate
9
+ Project-URL: Issues, https://github.com/Rasath16/GreenGate/issues
10
+ Keywords: green-ai,llm,carbon,sustainability,inference,model-cascade,routing,energy-efficiency,uncertainty
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3.10
17
+ Classifier: Programming Language :: Python :: 3.11
18
+ Classifier: Programming Language :: Python :: 3.12
19
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
20
+ Requires-Python: >=3.10
21
+ Description-Content-Type: text/markdown
22
+ License-File: LICENSE
23
+ Requires-Dist: torch>=2.0
24
+ Requires-Dist: transformers>=4.40
25
+ Provides-Extra: gpu
26
+ Requires-Dist: pynvml>=11.5; extra == "gpu"
27
+ Requires-Dist: bitsandbytes>=0.43; extra == "gpu"
28
+ Provides-Extra: api
29
+ Requires-Dist: ecologits; extra == "api"
30
+ Requires-Dist: openai; extra == "api"
31
+ Provides-Extra: vision
32
+ Requires-Dist: torchvision; extra == "vision"
33
+ Requires-Dist: pillow; extra == "vision"
34
+ Provides-Extra: eval
35
+ Requires-Dist: pandas; extra == "eval"
36
+ Requires-Dist: numpy; extra == "eval"
37
+ Requires-Dist: matplotlib; extra == "eval"
38
+ Requires-Dist: datasets; extra == "eval"
39
+ Dynamic: license-file
40
+
41
+ # GreenGate
42
+
43
+ **Cut the cost and carbon of LLM inference with one line of routing.** Easy queries are answered by a small local model; only genuinely hard ones escalate to a large model. Every query is carbon-accounted, honestly.
44
+
45
+ ```bash
46
+ pip install greengate
47
+ ```
48
+
49
+ ```python
50
+ import greengate
51
+
52
+ gw = greengate.GreenGate(
53
+ small="Qwen/Qwen2.5-0.5B-Instruct", # runs locally, exposes logits
54
+ large="gpt-4o-mini", # or any local model
55
+ budget_g=0.5, # optional carbon ceiling
56
+ )
57
+
58
+ r = gw.route("Summarise this document ...")
59
+ print(r.response) # answered by whichever tier was appropriate
60
+ print(r.decision) # "LOCAL" or "ESCALATE"
61
+ print(r.carbon_g) # gCO2 for this query, including any wasted small run
62
+
63
+ gw.profile() # session totals: carbon, escalation rate, latency
64
+ ```
65
+
66
+ ## Why this exists
67
+
68
+ Today every query, easy or hard, is sent to the same large model. Most queries do not need it. GreenGate sits between your application and your models, measures how uncertain the small model is, and escalates only when that uncertainty is high.
69
+
70
+ Three things it does that other cascading systems do not:
71
+
72
+ 1. **Measures real energy.** Local inference is metered with NVML at 100 ms across all GPUs, not estimated. API tiers use [EcoLogits](https://ecologits.ai/) and are labelled as estimates.
73
+ 2. **Full carbon accounting.** When a query escalates, it is charged for *both* the discarded small-model run and the large-model run. Most published cascade savings omit the first, which overstates them.
74
+ 3. **Routes vision queries too**, using the average token probability of the generated answer as the confidence signal.
75
+
76
+ ## What to expect
77
+
78
+ From the evaluation in the accompanying study (ShareGPT, MMLU and VQAv2; measured on dual T4):
79
+
80
+ | Workload | Carbon reduction | Quality retention |
81
+ |---|---|---|
82
+ | MMLU (structured) | 29.5% | 95.0% |
83
+ | ShareGPT (open-ended) | 49.8% | 80.6% |
84
+ | ShareGPT (quality-first) | 14.5% | 90.1% |
85
+ | VQAv2 (vision) | escalates only 5% | exceeds both tiers |
86
+
87
+ Savings depend on four measurable things: the energy ratio between your tiers, your escalation rate, the retention you accept, and your hardware. GreenGate reports all of them rather than assuming them. On an inefficient local GPU, escalating to an efficient API can genuinely be greener — the profiler tells you which case you are in.
88
+
89
+ ## Tiers
90
+
91
+ - **Small tier** must be a local open-weight model. The routing signal is computed from token logits, which hosted APIs do not expose. This is also where the privacy and cost win comes from.
92
+ - **Large tier** can be anything: another local model, or an OpenAI API model.
93
+
94
+ ## Calibration
95
+
96
+ Language models are overconfident, so raw entropy thresholds are unreliable. GreenGate ships fitted temperature-scaling values for evaluated models and self-calibrates for anything else:
97
+
98
+ ```python
99
+ gw.calibrate() # ~15 min, fits T on held-out MMLU validation, saved to ~/.greengate/
100
+ ```
101
+
102
+ ## Modes and budgets
103
+
104
+ ```python
105
+ gw.config(mode="green") # escalate less; "balanced" and "quality" also available
106
+ gw.config(threshold=2.9) # or set the entropy threshold explicitly
107
+ gw.config(budget_g=0.05) # sliding-window carbon ceiling; escalation defers when exhausted
108
+ ```
109
+
110
+ ## Honest limitations
111
+
112
+ - On open-ended generation with small models, token entropy is a weak signal (near chance in our evaluation). It is informative on structured tasks and vision. Where it is weak, savings come from the cascade structure rather than from selective routing.
113
+ - Energy measurement requires an NVIDIA GPU (NVML). CPU runs fall back to a documented estimate.
114
+ - API-tier carbon is an estimate, not a measurement, and is not directly comparable to metered local figures.
115
+
116
+ ## Install extras
117
+
118
+ ```bash
119
+ pip install greengate[gpu] # bitsandbytes + pynvml for quantised local models and metering
120
+ pip install greengate[api] # openai + ecologits for API large tiers
121
+ pip install greengate[eval] # pandas/matplotlib for the evaluation scripts
122
+ ```
123
+
124
+ ## Reproducing the evaluation
125
+
126
+ The `RUNBOOK.md` in this repository reproduces every published number: calibration, three deployment configurations, four baselines, threshold sweeps, grid conditions, trace replay against real Azure arrival traces, and three ablation studies. Raw per-query records for all runs are in `experiments/`.
127
+
128
+ ## Citation
129
+
130
+ If you use GreenGate in academic work, please cite the accompanying study:
131
+
132
+ > T. R. Hemachandra, "GreenGate: A Confidence-Aware Cascading Framework for Optimizing Energy and Cost in Large Language and Multimodal Model Inference," BSc thesis, NSBM Green University, 2026.
133
+
134
+ ## License
135
+
136
+ MIT
@@ -0,0 +1,96 @@
1
+ # GreenGate
2
+
3
+ **Cut the cost and carbon of LLM inference with one line of routing.** Easy queries are answered by a small local model; only genuinely hard ones escalate to a large model. Every query is carbon-accounted, honestly.
4
+
5
+ ```bash
6
+ pip install greengate
7
+ ```
8
+
9
+ ```python
10
+ import greengate
11
+
12
+ gw = greengate.GreenGate(
13
+ small="Qwen/Qwen2.5-0.5B-Instruct", # runs locally, exposes logits
14
+ large="gpt-4o-mini", # or any local model
15
+ budget_g=0.5, # optional carbon ceiling
16
+ )
17
+
18
+ r = gw.route("Summarise this document ...")
19
+ print(r.response) # answered by whichever tier was appropriate
20
+ print(r.decision) # "LOCAL" or "ESCALATE"
21
+ print(r.carbon_g) # gCO2 for this query, including any wasted small run
22
+
23
+ gw.profile() # session totals: carbon, escalation rate, latency
24
+ ```
25
+
26
+ ## Why this exists
27
+
28
+ Today every query, easy or hard, is sent to the same large model. Most queries do not need it. GreenGate sits between your application and your models, measures how uncertain the small model is, and escalates only when that uncertainty is high.
29
+
30
+ Three things it does that other cascading systems do not:
31
+
32
+ 1. **Measures real energy.** Local inference is metered with NVML at 100 ms across all GPUs, not estimated. API tiers use [EcoLogits](https://ecologits.ai/) and are labelled as estimates.
33
+ 2. **Full carbon accounting.** When a query escalates, it is charged for *both* the discarded small-model run and the large-model run. Most published cascade savings omit the first, which overstates them.
34
+ 3. **Routes vision queries too**, using the average token probability of the generated answer as the confidence signal.
35
+
36
+ ## What to expect
37
+
38
+ From the evaluation in the accompanying study (ShareGPT, MMLU and VQAv2; measured on dual T4):
39
+
40
+ | Workload | Carbon reduction | Quality retention |
41
+ |---|---|---|
42
+ | MMLU (structured) | 29.5% | 95.0% |
43
+ | ShareGPT (open-ended) | 49.8% | 80.6% |
44
+ | ShareGPT (quality-first) | 14.5% | 90.1% |
45
+ | VQAv2 (vision) | escalates only 5% | exceeds both tiers |
46
+
47
+ Savings depend on four measurable things: the energy ratio between your tiers, your escalation rate, the retention you accept, and your hardware. GreenGate reports all of them rather than assuming them. On an inefficient local GPU, escalating to an efficient API can genuinely be greener — the profiler tells you which case you are in.
48
+
49
+ ## Tiers
50
+
51
+ - **Small tier** must be a local open-weight model. The routing signal is computed from token logits, which hosted APIs do not expose. This is also where the privacy and cost win comes from.
52
+ - **Large tier** can be anything: another local model, or an OpenAI API model.
53
+
54
+ ## Calibration
55
+
56
+ Language models are overconfident, so raw entropy thresholds are unreliable. GreenGate ships fitted temperature-scaling values for evaluated models and self-calibrates for anything else:
57
+
58
+ ```python
59
+ gw.calibrate() # ~15 min, fits T on held-out MMLU validation, saved to ~/.greengate/
60
+ ```
61
+
62
+ ## Modes and budgets
63
+
64
+ ```python
65
+ gw.config(mode="green") # escalate less; "balanced" and "quality" also available
66
+ gw.config(threshold=2.9) # or set the entropy threshold explicitly
67
+ gw.config(budget_g=0.05) # sliding-window carbon ceiling; escalation defers when exhausted
68
+ ```
69
+
70
+ ## Honest limitations
71
+
72
+ - On open-ended generation with small models, token entropy is a weak signal (near chance in our evaluation). It is informative on structured tasks and vision. Where it is weak, savings come from the cascade structure rather than from selective routing.
73
+ - Energy measurement requires an NVIDIA GPU (NVML). CPU runs fall back to a documented estimate.
74
+ - API-tier carbon is an estimate, not a measurement, and is not directly comparable to metered local figures.
75
+
76
+ ## Install extras
77
+
78
+ ```bash
79
+ pip install greengate[gpu] # bitsandbytes + pynvml for quantised local models and metering
80
+ pip install greengate[api] # openai + ecologits for API large tiers
81
+ pip install greengate[eval] # pandas/matplotlib for the evaluation scripts
82
+ ```
83
+
84
+ ## Reproducing the evaluation
85
+
86
+ The `RUNBOOK.md` in this repository reproduces every published number: calibration, three deployment configurations, four baselines, threshold sweeps, grid conditions, trace replay against real Azure arrival traces, and three ablation studies. Raw per-query records for all runs are in `experiments/`.
87
+
88
+ ## Citation
89
+
90
+ If you use GreenGate in academic work, please cite the accompanying study:
91
+
92
+ > T. R. Hemachandra, "GreenGate: A Confidence-Aware Cascading Framework for Optimizing Energy and Cost in Large Language and Multimodal Model Inference," BSc thesis, NSBM Green University, 2026.
93
+
94
+ ## License
95
+
96
+ MIT
@@ -0,0 +1,12 @@
1
+ from greengate.core import GreenGate, RouteResult
2
+
3
+ __version__ = "0.1.0"
4
+ __all__ = ["GreenGate", "RouteResult"]
5
+
6
+
7
+ def __getattr__(name):
8
+ # legacy demo class, loaded lazily to keep `import greengate` light
9
+ if name == "GreenGateRouter":
10
+ from greengate.router import GreenGateRouter
11
+ return GreenGateRouter
12
+ raise AttributeError(name)
@@ -0,0 +1,153 @@
1
+ """Large-tier inference via OpenAI API with EcoLogits carbon estimation.
2
+
3
+ Dual-Mode Carbon Profiler, API side: local models are measured with
4
+ pynvml (profiler.py); API models are estimated with EcoLogits
5
+ (Rince et al., 2025, JOSS). If EcoLogits is unavailable or fails,
6
+ falls back to a documented per-token estimate and flags it.
7
+ """
8
+
9
+ import os
10
+ import time
11
+ from dataclasses import dataclass
12
+
13
+ # Fallback constants (used ONLY if EcoLogits fails) — documented in thesis:
14
+ # Jegham et al. (2025) measure ~0.42 Wh for a mean GPT-4o query (~150 output
15
+ # tokens). GPT-4o-mini is a distilled/smaller deployment; we conservatively
16
+ # assume 0.5x the 4o figure, scaled linearly by output tokens.
17
+ FALLBACK_WH_PER_OUTPUT_TOKEN = 0.42 * 0.5 / 150 # = 0.0014 Wh/token
18
+ WORLD_CARBON_INTENSITY = 475.0 # gCO2/kWh
19
+ DATACENTER_PUE = 1.2
20
+
21
+
22
+ @dataclass
23
+ class APIResult:
24
+ response: str
25
+ model: str
26
+ input_tokens: int
27
+ output_tokens: int
28
+ latency_s: float
29
+ energy_wh: float
30
+ carbon_grams: float
31
+ carbon_source: str # "ecologits" or "fallback_estimate"
32
+ # confidence signals from top-20 logprobs (None unless logprobs=True):
33
+ avg_logprob: float | None = None # mean log P(chosen token)
34
+ trunc_entropy: float | None = None # mean entropy over renormalised top-20 (bits)
35
+
36
+
37
+ class APILargeTier:
38
+ def __init__(self, model: str = "gpt-4o-mini", max_tokens: int = 300,
39
+ dry_run: bool = False, logprobs: bool = False):
40
+ self.model = model
41
+ self.max_tokens = max_tokens
42
+ self.dry_run = dry_run
43
+ self.logprobs = logprobs
44
+ self._ecologits_ok = False
45
+
46
+ if dry_run:
47
+ return
48
+
49
+ # EcoLogits must be initialised BEFORE the OpenAI client is created.
50
+ try:
51
+ from ecologits import EcoLogits
52
+ EcoLogits.init(providers=["openai"])
53
+ self._ecologits_ok = True
54
+ except Exception as e:
55
+ print(f" [api_tier] EcoLogits unavailable ({e}) — using documented fallback estimate")
56
+
57
+ from openai import OpenAI
58
+ self.client = OpenAI() # reads OPENAI_API_KEY env var
59
+
60
+ def query_vision(self, prompt: str, image) -> APIResult:
61
+ """Vision escalation: PIL image sent as a base64 data URL."""
62
+ if self.dry_run:
63
+ return self.query(prompt)
64
+ import base64
65
+ import io
66
+ buf = io.BytesIO()
67
+ image.save(buf, format="JPEG", quality=85)
68
+ b64 = base64.b64encode(buf.getvalue()).decode()
69
+ content = [
70
+ {"type": "text", "text": prompt},
71
+ {"type": "image_url",
72
+ "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
73
+ ]
74
+ return self._chat(content)
75
+
76
+ def query(self, prompt: str) -> APIResult:
77
+ if self.dry_run:
78
+ return APIResult(
79
+ response="[DRY RUN — no API call made]",
80
+ model=self.model, input_tokens=len(prompt) // 4,
81
+ output_tokens=50, latency_s=0.5,
82
+ energy_wh=50 * FALLBACK_WH_PER_OUTPUT_TOKEN,
83
+ carbon_grams=50 * FALLBACK_WH_PER_OUTPUT_TOKEN / 1000
84
+ * DATACENTER_PUE * WORLD_CARBON_INTENSITY,
85
+ carbon_source="dry_run",
86
+ )
87
+
88
+ return self._chat(prompt)
89
+
90
+ def _chat(self, content, retries: int = 6) -> APIResult:
91
+ kwargs = {}
92
+ if self.logprobs:
93
+ kwargs = {"logprobs": True, "top_logprobs": 20}
94
+ from openai import RateLimitError, APIError
95
+ t0 = time.perf_counter()
96
+ for attempt in range(retries):
97
+ try:
98
+ resp = self.client.chat.completions.create(
99
+ model=self.model,
100
+ messages=[{"role": "user", "content": content}],
101
+ max_tokens=self.max_tokens,
102
+ **kwargs,
103
+ )
104
+ break
105
+ except (RateLimitError, APIError):
106
+ if attempt == retries - 1:
107
+ raise
108
+ time.sleep(min(2 ** attempt, 30)) # 1,2,4,8,16,30s backoff
109
+ latency = time.perf_counter() - t0
110
+
111
+ avg_lp = trunc_H = None
112
+ lp_content = getattr(resp.choices[0].logprobs, "content", None) \
113
+ if self.logprobs and resp.choices[0].logprobs else None
114
+ if lp_content:
115
+ import math
116
+ lps, ents = [], []
117
+ for tok in lp_content:
118
+ lps.append(tok.logprob)
119
+ if tok.top_logprobs:
120
+ ps = [math.exp(t.logprob) for t in tok.top_logprobs]
121
+ z = sum(ps) or 1.0
122
+ ents.append(-sum(p / z * math.log2(p / z) for p in ps if p > 0))
123
+ avg_lp = sum(lps) / len(lps) if lps else None
124
+ trunc_H = sum(ents) / len(ents) if ents else None
125
+
126
+ text = resp.choices[0].message.content or ""
127
+ in_tok = resp.usage.prompt_tokens
128
+ out_tok = resp.usage.completion_tokens
129
+
130
+ energy_wh, carbon_g, source = None, None, "fallback_estimate"
131
+ if self._ecologits_ok and hasattr(resp, "impacts"):
132
+ try:
133
+ # EcoLogits attaches .impacts; energy in kWh (value or range)
134
+ e = resp.impacts.energy.value
135
+ energy_kwh = (e.min + e.max) / 2 if hasattr(e, "min") else float(e)
136
+ g = resp.impacts.gwp.value
137
+ gwp_kg = (g.min + g.max) / 2 if hasattr(g, "min") else float(g)
138
+ energy_wh = energy_kwh * 1000
139
+ carbon_g = gwp_kg * 1000
140
+ source = "ecologits"
141
+ except Exception:
142
+ pass
143
+ if energy_wh is None:
144
+ energy_wh = out_tok * FALLBACK_WH_PER_OUTPUT_TOKEN
145
+ carbon_g = energy_wh / 1000 * DATACENTER_PUE * WORLD_CARBON_INTENSITY
146
+
147
+ return APIResult(
148
+ response=text.strip(), model=self.model,
149
+ input_tokens=in_tok, output_tokens=out_tok,
150
+ latency_s=latency, energy_wh=energy_wh,
151
+ carbon_grams=carbon_g, carbon_source=source,
152
+ avg_logprob=avg_lp, trunc_entropy=trunc_H,
153
+ )
@@ -0,0 +1,31 @@
1
+ """Sliding-window carbon budget (Chapter 3, Design Analysis 2.5.3).
2
+
3
+ Maintains a window of the last `window_s` seconds of routing decisions.
4
+ If the carbon spent inside the window would exceed budget K, escalations
5
+ are temporarily blocked (the query falls back to the small answer) until
6
+ older queries age out of the window.
7
+ """
8
+
9
+ from collections import deque
10
+
11
+
12
+ class SlidingWindowBudget:
13
+ def __init__(self, budget_g: float, window_s: float = 3600.0):
14
+ self.budget_g = budget_g
15
+ self.window_s = window_s
16
+ self._events: deque[tuple[float, float]] = deque() # (timestamp, carbon_g)
17
+
18
+ def _expire(self, now: float):
19
+ while self._events and self._events[0][0] < now - self.window_s:
20
+ self._events.popleft()
21
+
22
+ def window_carbon(self, now: float) -> float:
23
+ self._expire(now)
24
+ return sum(c for _, c in self._events)
25
+
26
+ def allows(self, now: float, escalation_cost_g: float) -> bool:
27
+ """Would escalating now keep the window under budget?"""
28
+ return self.window_carbon(now) + escalation_cost_g <= self.budget_g
29
+
30
+ def record(self, now: float, carbon_g: float):
31
+ self._events.append((now, carbon_g))
@@ -0,0 +1,53 @@
1
+ """Temperature scaling calibration (Guo et al., 2017).
2
+
3
+ Fits a single scalar T on a held-out validation split by minimising
4
+ negative log-likelihood. T > 1 softens overconfident distributions.
5
+ Reports Expected Calibration Error (ECE) before and after.
6
+ """
7
+
8
+ import torch
9
+ import torch.nn.functional as F
10
+
11
+
12
+ def fit_temperature(logits: torch.Tensor, labels: torch.Tensor,
13
+ t_min: float = 0.25, t_max: float = 8.0,
14
+ n_grid: int = 400) -> float:
15
+ """Fit T by fine grid search over NLL (robust, dependency-free).
16
+
17
+ Args:
18
+ logits: (n, k) raw choice logits from the validation split
19
+ labels: (n,) correct choice indices
20
+ """
21
+ best_t, best_nll = 1.0, float("inf")
22
+ for i in range(n_grid):
23
+ t = t_min + (t_max - t_min) * i / (n_grid - 1)
24
+ nll = F.cross_entropy(logits / t, labels).item()
25
+ if nll < best_nll:
26
+ best_nll, best_t = nll, t
27
+ return best_t
28
+
29
+
30
+ def ece(logits: torch.Tensor, labels: torch.Tensor,
31
+ temperature: float = 1.0, n_bins: int = 15) -> float:
32
+ """Expected Calibration Error over equal-width confidence bins."""
33
+ probs = F.softmax(logits / temperature, dim=-1)
34
+ conf, pred = probs.max(dim=-1)
35
+ correct = (pred == labels).float()
36
+
37
+ total = 0.0
38
+ n = len(labels)
39
+ for b in range(n_bins):
40
+ lo, hi = b / n_bins, (b + 1) / n_bins
41
+ mask = (conf > lo) & (conf <= hi)
42
+ if mask.sum() == 0:
43
+ continue
44
+ avg_conf = conf[mask].mean().item()
45
+ avg_acc = correct[mask].mean().item()
46
+ total += (mask.sum().item() / n) * abs(avg_conf - avg_acc)
47
+ return total
48
+
49
+
50
+ def calibrated_entropy(logits: torch.Tensor, temperature: float) -> float:
51
+ """Shannon entropy (bits) of the temperature-scaled distribution."""
52
+ probs = F.softmax(logits / temperature, dim=-1).clamp(min=1e-9)
53
+ return -(probs * probs.log2()).sum(dim=-1).item()