Модуль 6. Кодирование источника: от энтропии к реальным битам

pip install torch transformers zstandard

from __future__ import annotations

import gzip
import hashlib
import math
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

try:
    import zstandard as zstd
except ImportError as exc:
    raise SystemExit("Install zstandard: pip install zstandard") from exc


FILE_PATH = Path("corpus.txt")
MODEL_NAME = "openai-community/gpt2"
MODEL_REVISION = "607a30d783dfa663caf39e06633721c8d4cfcd7e"
ZSTD_LEVEL = 19
REQUESTED_STRIDE = 512


def read_utf8_file(path: Path) -> tuple[bytes, str]:
    raw = path.read_bytes()
    text = raw.decode("utf-8")
    if not raw:
        raise ValueError("The input file is empty.")
    return raw, text


def model_max_length(model: torch.nn.Module) -> int:
    for name in ("max_position_embeddings", "n_positions", "n_ctx"):
        value = getattr(model.config, name, None)
        if isinstance(value, int) and value > 1:
            return value
    raise ValueError("Could not determine the model context length.")


def causal_lm_nll(
    text: str,
    model_name: str,
    requested_stride: int = 512,
    revision: str = MODEL_REVISION,
) -> dict[str, float | int | str]:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    tokenizer = AutoTokenizer.from_pretrained(model_name, revision=revision, use_fast=True)
    model = AutoModelForCausalLM.from_pretrained(model_name, revision=revision).to(device).eval()

    token_ids = tokenizer(
        text,
        add_special_tokens=False,
        return_tensors="pt",
    ).input_ids

    n_text_tokens = int(token_ids.size(1))
    if n_text_tokens == 0:
        raise ValueError("The tokenizer produced no tokens.")

    # GPT-2 has no separate BOS token. Using EOS as the shared start symbol
    # defines the probability assigned to the first text token.
    start_id = tokenizer.bos_token_id
    if start_id is None:
        start_id = tokenizer.eos_token_id
    if start_id is None:
        raise ValueError("The tokenizer has neither BOS nor EOS token.")

    start = torch.tensor([[start_id]], dtype=token_ids.dtype)
    ids = torch.cat([start, token_ids], dim=1)

    max_len = model_max_length(model)
    stride = min(requested_stride, max_len - 1)
    if stride < 1:
        raise ValueError("The context length must be at least 2.")

    nll_sum_nats = 0.0
    n_counted = 0

    # Positions 1..N in `ids` are the original text tokens. Each target
    # interval is disjoint, while its input window may overlap for context.
    target_start = 1
    while target_start < ids.size(1):
        target_end = min(target_start + stride, ids.size(1))
        begin = max(0, target_end - max_len)

        chunk = ids[:, begin:target_end].to(device)
        labels = chunk.clone()
        local_target_start = target_start - begin
        labels[:, :local_target_start] = -100

        with torch.inference_mode():
            output = model(input_ids=chunk, labels=labels)

        # Causal-LM loss compares logits[:, :-1] with labels[:, 1:].
        # Counting after this shift avoids the common off-by-one bug.
        n_loss_tokens = int((labels[:, 1:] != -100).sum().item())
        expected = target_end - target_start
        if n_loss_tokens != expected:
            raise RuntimeError(
                f"Window accounting error: {n_loss_tokens=} != {expected=}"
            )

        nll_sum_nats += float(output.loss.item()) * n_loss_tokens
        n_counted += n_loss_tokens
        target_start = target_end

    if n_counted != n_text_tokens:
        raise RuntimeError(
            f"Counted {n_counted} targets, expected {n_text_tokens}."
        )

    total_bits = nll_sum_nats / math.log(2.0)
    return {
        "device": str(device),
        "n_tokens": n_text_tokens,
        "nll_nats": nll_sum_nats,
        "ppl": math.exp(nll_sum_nats / n_counted),
        "bits_per_token": total_bits / n_counted,
        "total_bits": total_bits,
    }


def main() -> None:
    raw, text = read_utf8_file(FILE_PATH)

    gzip_bytes = gzip.compress(raw, compresslevel=9)
    zstd_bytes = zstd.ZstdCompressor(level=ZSTD_LEVEL).compress(raw)
    lm = causal_lm_nll(text, MODEL_NAME, REQUESTED_STRIDE)

    n_bytes = len(raw)
    n_chars = len(text)  # Python Unicode code points, not grapheme clusters.
    theoretical_bytes = float(lm["total_bits"]) / 8.0

    print(f"File: {FILE_PATH}")
    print(f"UTF-8 bytes: {n_bytes:,}; Python characters: {n_chars:,}")
    print(f"Model: {MODEL_NAME}; device: {lm['device']}")
    print(f"Model revision: {MODEL_REVISION}")
    print(f"Corpus SHA-256: {hashlib.sha256(raw).hexdigest()}")
    print(f"Counted text tokens: {lm['n_tokens']:,}")
    print()

    print("Real archives (headers included):")
    print(
        f"  gzip -9: {len(gzip_bytes):,} bytes, "
        f"{8 * len(gzip_bytes) / n_bytes:.4f} BPB"
    )
    print(
        f"  zstd -{ZSTD_LEVEL}: {len(zstd_bytes):,} bytes, "
        f"{8 * len(zstd_bytes) / n_bytes:.4f} BPB"
    )
    print()

    print("Causal-LM NLL-equivalent code (not an emitted archive):")
    print(f"  perplexity: {lm['ppl']:.4f}")
    print(f"  bits/token: {lm['bits_per_token']:.4f}")
    print(f"  theoretical bytes: {theoretical_bytes:,.1f}")
    print(f"  BPB: {float(lm['total_bits']) / n_bytes:.4f}")
    print(f"  BPC: {float(lm['total_bits']) / n_chars:.4f}")
    print()
    print("Model weights, tokenizer files, and entropy-coder overhead are excluded.")


if __name__ == "__main__":
    main()

n_loss_tokens = int((labels[:, 1:] != -100).sum().item())

assert n_counted == n_text_tokens

from __future__ import annotations

import math
import random


def log2sumexp(values: list[float]) -> float:
    m = max(values)
    return m + math.log2(sum(2.0 ** (v - m) for v in values))


def bernoulli_log2_prob(sequence: list[int], p: float) -> float:
    ones = sum(sequence)
    zeros = len(sequence) - ones
    return ones * math.log2(p) + zeros * math.log2(1.0 - p)


def check_sequence(sequence: list[int]) -> tuple[float, float]:
    parameters = [0.2, 0.8]
    log2_weights = [math.log2(0.5), math.log2(0.5)]
    model_log_probs = [
        bernoulli_log2_prob(sequence, p) for p in parameters
    ]

    mixture_log_prob = log2sumexp(
        [w + lp for w, lp in zip(log2_weights, model_log_probs)]
    )
    mixture_length = -mixture_log_prob
    best_model_length = min(-lp for lp in model_log_probs)

    if mixture_length > best_model_length + 1.0 + 1e-12:
        raise AssertionError("Mixture regret bound was violated.")

    return mixture_length, best_model_length


rng = random.Random(6)
for n in (1, 5, 20, 100, 1000):
    for _ in range(100):
        sequence = [rng.randrange(2) for _ in range(n)]
        mixture_length, best_length = check_sequence(sequence)
    print(
        f"n={n:4d}: last mixture={mixture_length:10.4f} bits, "
        f"best={best_length:10.4f}, regret={mixture_length-best_length:.4f}"
    )