Module 2. Entropy

from __future__ import annotations

import hashlib
import math
import re
from collections import Counter
from dataclasses import dataclass

import requests

URL = "https://www.gutenberg.org/files/11/11-0.txt"
ALPHABET = "abcdefghijklmnopqrstuvwxyz "
ALPHABET_SIZE = len(ALPHABET)
REFERENCE_TEXT_SHA256 = "c7172a25460df6168670a3ac6c778482feb23546fabf24d05ba9d19dbb9256c8"


def download_and_normalize(url: str = URL) -> str:
    """Download Project Gutenberg text and map it to a 27-symbol alphabet."""
    response = requests.get(url, timeout=30)
    response.raise_for_status()
    raw = response.content.decode("utf-8")

    # Remove the Project Gutenberg wrapper while tolerating minor header changes.
    start = re.search(
        r"\*\*\*\s*START OF THE PROJECT GUTENBERG EBOOK.*?\*\*\*",
        raw,
        flags=re.IGNORECASE,
    )
    end = re.search(
        r"\*\*\*\s*END OF THE PROJECT GUTENBERG EBOOK.*?\*\*\*",
        raw,
        flags=re.IGNORECASE,
    )
    if start is not None:
        raw = raw[start.end() :]
    if end is not None:
        # If the start wrapper was removed, search again in the shortened string.
        end_after_start = re.search(
            r"\*\*\*\s*END OF THE PROJECT GUTENBERG EBOOK.*?\*\*\*",
            raw,
            flags=re.IGNORECASE,
        )
        if end_after_start is not None:
            raw = raw[: end_after_start.start()]

    text = re.sub(r"[^a-z]+", " ", raw.lower())
    text = re.sub(r" +", " ", text).strip()
    if not text:
        raise ValueError("Normalization produced an empty text")
    digest = hashlib.sha256(text.encode("utf-8")).hexdigest()
    print(f"Normalized corpus SHA-256: {digest}")
    if digest != REFERENCE_TEXT_SHA256:
        raise ValueError("Corpus differs from the reference input; compare the text before comparing tables")
    return text


def plugin_conditional_entropy(text: str, k: int) -> float:
    """Empirical H(next symbol | context of length k), evaluated in-sample."""
    if not 0 <= k < len(text):
        raise ValueError("k must satisfy 0 <= k < len(text)")

    pair_counts: Counter[tuple[str, str]] = Counter()
    context_counts: Counter[str] = Counter()

    for i in range(k, len(text)):
        context = text[i - k : i] if k else ""
        symbol = text[i]
        pair_counts[(context, symbol)] += 1
        context_counts[context] += 1

    total = len(text) - k
    return -sum(
        (count / total) * math.log2(count / context_counts[context])
        for (context, _symbol), count in pair_counts.items()
    )


@dataclass(frozen=True)
class NGramCounts:
    pair: Counter[tuple[str, str]]
    context: Counter[str]


def fit_ngram(text: str, k: int) -> NGramCounts:
    """Count context/next-symbol pairs in the training prefix."""
    pair_counts: Counter[tuple[str, str]] = Counter()
    context_counts: Counter[str] = Counter()

    for i in range(k, len(text)):
        context = text[i - k : i] if k else ""
        symbol = text[i]
        pair_counts[(context, symbol)] += 1
        context_counts[context] += 1

    return NGramCounts(pair=pair_counts, context=context_counts)


def heldout_cross_entropy(
    full_text: str,
    split: int,
    k: int,
    alpha: float = 0.1,
) -> float:
    """Score test targets using an add-alpha model fitted only on the prefix."""
    if not 0 < split < len(full_text):
        raise ValueError("split must lie strictly inside the text")
    if not 0 <= k <= split:
        raise ValueError("k must satisfy 0 <= k <= split")
    if alpha <= 0:
        raise ValueError("alpha must be positive")

    train = full_text[:split]
    counts = fit_ngram(train, k)

    nll_bits = 0.0
    n_targets = 0

    # Every scored target lies in test. For the first k test targets, the context
    # may cross the train/test boundary; this is legitimate chronological context.
    for i in range(split, len(full_text)):
        context = full_text[i - k : i] if k else ""
        symbol = full_text[i]
        numerator = counts.pair[(context, symbol)] + alpha
        denominator = counts.context[context] + alpha * ALPHABET_SIZE
        probability = numerator / denominator
        nll_bits -= math.log2(probability)
        n_targets += 1

    if n_targets == 0:
        raise ValueError("The test split contains no targets")
    return nll_bits / n_targets


text = download_and_normalize()
print(f"Normalized length: {len(text):,} characters")

print("\nIn-sample plug-in conditional entropy")
for k in range(10):
    value = plugin_conditional_entropy(text, k)
    print(f"k={k:2d}: {value:.4f} bits/character")

split = int(0.8 * len(text))
print("\nHeld-out add-0.1 cross-entropy")
for k in range(8):
    value = heldout_cross_entropy(text, split, k, alpha=0.1)
    print(f"k={k:2d}: {value:.4f} bits/character")

import math

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "openai-community/gpt2"
MODEL_REVISION = "607a30d783dfa663caf39e06633721c8d4cfcd7e"
SAMPLE_CHARS = 50_000
STRIDE = 512

# `text` is the normalized string from Exercise 4.
sample = text[:SAMPLE_CHARS]
if not sample:
    raise ValueError("The sample is empty")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=MODEL_REVISION)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, revision=MODEL_REVISION).to(device).eval()

text_ids = tokenizer(
    sample,
    return_tensors="pt",
    add_special_tokens=False,
).input_ids

if tokenizer.eos_token_id is None:
    raise ValueError("The tokenizer has no EOS token for the initial prefix")

# Prefix EOS lets the model assign probability to the first text token as well.
eos = torch.tensor([[tokenizer.eos_token_id]], dtype=text_ids.dtype)
input_ids = torch.cat([eos, text_ids], dim=1).to(device)

max_length = int(model.config.n_positions)
if not 0 < STRIDE < max_length:
    raise ValueError("STRIDE must satisfy 0 < STRIDE < model context length")

nll_sum_nats = 0.0
n_scored_tokens = 0
previous_end = 0
sequence_length = input_ids.size(1)

for begin in range(0, sequence_length, STRIDE):
    end = min(begin + max_length, sequence_length)
    new_length = end - previous_end

    chunk = input_ids[:, begin:end]
    labels = chunk.clone()
    labels[:, :-new_length] = -100

    # Hugging Face causal-LM loss compares logits[:, :-1] with labels[:, 1:].
    # Count exactly those shifted labels that are not masked.
    n_loss_tokens = int((labels[:, 1:] != -100).sum().item())

    if n_loss_tokens > 0:
        with torch.no_grad():
            output = model(chunk, labels=labels)
        nll_sum_nats += float(output.loss.item()) * n_loss_tokens
        n_scored_tokens += n_loss_tokens

    previous_end = end
    if end == sequence_length:
        break

expected_scored = text_ids.size(1)
if n_scored_tokens != expected_scored:
    raise RuntimeError(
        f"Scored {n_scored_tokens} tokens, expected {expected_scored}"
    )

bits_per_token = nll_sum_nats / (n_scored_tokens * math.log(2))
bits_per_character = nll_sum_nats / (len(sample) * math.log(2))
perplexity = math.exp(nll_sum_nats / n_scored_tokens)

print(f"Scored tokens:      {n_scored_tokens:,}")
print(f"Bits per token:     {bits_per_token:.4f}")
print(f"Bits per character: {bits_per_character:.4f}")
print(f"Token perplexity:   {perplexity:.4f}")