| """Mean per-sample empirical unigram entropy in natural-log units.""" |
|
|
| from __future__ import annotations |
|
|
| from collections import Counter |
| from dataclasses import dataclass |
| import math |
| from typing import Iterable, Mapping, Sequence |
|
|
|
|
| @dataclass(frozen=True) |
| class EntropyResult: |
| mean_entropy: float |
| sample_count: int |
| token_count: int |
| excluded_token_ids: tuple[int, ...] |
| aggregation: str = "arithmetic_mean_of_per_sample_entropy" |
| logarithm: str = "natural" |
| unit: str = "nats" |
|
|
|
|
| def unigram_entropy(token_ids: Sequence[int]) -> float: |
| """Return the natural-log empirical entropy of one non-empty token row.""" |
|
|
| if not token_ids: |
| raise ValueError("token row has no tokens") |
| if not all(type(token_id) is int and token_id >= 0 for token_id in token_ids): |
| raise ValueError("token IDs must be non-negative integers") |
| total = len(token_ids) |
| return -sum( |
| (count / total) * math.log(count / total) for count in Counter(token_ids).values() |
| ) |
|
|
|
|
| def _token_ids(record: object, index: int) -> Sequence[int]: |
| if isinstance(record, Mapping): |
| value = record.get("token_ids") |
| else: |
| value = getattr(record, "token_ids", None) |
| if not isinstance(value, (list, tuple)): |
| raise ValueError(f"record {index} has invalid token_ids") |
| return value |
|
|
|
|
| def mean_unigram_entropy( |
| records: Iterable[object], special_ids: Iterable[int] = () |
| ) -> EntropyResult: |
| """Remove exactly the caller-declared padding IDs, then average row entropy.""" |
|
|
| excluded = frozenset(special_ids) |
| if not all(type(token_id) is int and token_id >= 0 for token_id in excluded): |
| raise ValueError("special token IDs must be non-negative integers") |
| entropies: list[float] = [] |
| token_count = 0 |
| for index, record in enumerate(records): |
| row = [token_id for token_id in _token_ids(record, index) if token_id not in excluded] |
| if not row: |
| raise ValueError(f"record {index} has no tokens after exclusions") |
| entropies.append(unigram_entropy(row)) |
| token_count += len(row) |
| if not entropies: |
| raise ValueError("no sample records were provided") |
| result = math.fsum(entropies) / len(entropies) |
| if not math.isfinite(result): |
| raise ValueError("mean entropy is not finite") |
| return EntropyResult( |
| mean_entropy=result, |
| sample_count=len(entropies), |
| token_count=token_count, |
| excluded_token_ids=tuple(sorted(excluded)), |
| ) |
|
|
