""" ============================================================================= 第 11 讲 · 代码 1：预测就是压缩——用第 1 讲的模型做真正的算术编码 ============================================================================= 对应讲义：lectures/lecture11_compression.md 第 2、3 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- "一个好的语言模型就是一个好的压缩器"——这句话可以逐字节地验证。 演示 1：香农码长 −log2 p；为什么要算术编码：哈夫曼编码每个符号至少 1 bit，概率很偏时浪费很多。 演示 2：把第 1 讲的验证集文本（14,085 个字符）分别用均匀分布、n-gram、base model 做算术编码， 和 gzip / bz2 / lzma 比较；解压回来，逐字检查。 演示 3："边解边预测"：解码器事先不知道文本，每解出一个字符才把它交给模型——真的这样跑一遍。 演示 4：模型自己也要算钱：base model 把它的训练文本压得极小，但模型本身比数据还大。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 01_arithmetic_coding.py # 约 10 秒；需要 PyTorch 依赖本目录的 codec_tools.py，以及 code/lecture01 的 fm_world.py、ngram_lm.py、tiny_gpt.py、console.py （第 1 讲的 base model 只读加载）。 【Python 小知识】 - gzip / bz2 / lzma 都是 Python 标准库里的通用压缩器：xxx.compress(字节串) 返回压缩后的字节串。 它们不知道"英语"或"迷你互联网"，只会找重复的片段（gzip、lzma）或做块排序变换（bz2）。 - heapq：最小堆，heappush / heappop 每次取出最小的元素——构造哈夫曼树用。 """ import bz2 import gzip import heapq import lzma import math import os import random import sys import torch HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from codec_tools import (ArithmeticDecoder, ArithmeticEncoder, decode_text, encode_text, gpt_distributions, # noqa: E402 ideal_bits, ngram_distributions, quantize) from console import pad, print_header # noqa: E402 from fm_world import ALPHABET, build_world # noqa: E402 from ngram_lm import CharNGramLM # noqa: E402 from tiny_gpt import CharTokenizer, count_parameters, load_or_train_base_model # noqa: E402 # ============================================================================= # 演示 1：香农码长、哈夫曼与算术编码 # ============================================================================= def huffman_lengths(probs): """ 哈夫曼编码：每个符号的码字长度（bit）。 【思路】反复把概率最小的两个节点合并成一个新节点（概率相加），直到只剩一个； 每个符号的码长 = 它从叶子到根经过了几次合并。哈夫曼编码是"每个符号用整数个 bit"时最优的前缀码。 【参数】probs：概率列表 【返回】list[int]，每个符号的码长 """ heap = [(p, i, [i]) for i, p in enumerate(probs)] # (概率, 用来打破平局的编号, 这个节点包含的符号) heapq.heapify(heap) lengths = [0] * len(probs) counter = len(probs) while len(heap) > 1: p1, _, s1 = heapq.heappop(heap) p2, _, s2 = heapq.heappop(heap) for s in s1 + s2: lengths[s] += 1 heapq.heappush(heap, (p1 + p2, counter, s1 + s2)) counter += 1 return lengths def demo1_codes(): print_header("演示 1：一个符号要几个 bit——熵、哈夫曼编码、算术编码") rng = random.Random(0) sources = [("A：概率是 2 的幂", [0.5, 0.25, 0.125, 0.125]), ("B：很偏的分布", [0.9, 0.05, 0.03, 0.02])] print(f"{pad('信源', 18)}{pad('熵 H（bit/符号）', 18, '>')}{pad('哈夫曼 平均码长', 18, '>')}{pad('算术编码 实测', 16, '>')}") for name, p in sources: H = -sum(q * math.log2(q) for q in p) L = huffman_lengths(p) huff = sum(q * l for q, l in zip(p, L)) symbols = rng.choices(range(len(p)), weights=p, k=10000) enc = ArithmeticEncoder() freqs = quantize(p) for s in symbols: enc.encode(freqs, s) bits = enc.finish() dec = ArithmeticDecoder(bits) assert [dec.decode(freqs) for _ in symbols] == symbols print(f"{pad(name, 18)}{pad(f'{H:.4f}', 18, '>')}{pad(f'{huff:.4f}', 18, '>')}{pad(f'{len(bits) / len(symbols):.4f}', 16, '>')}") print(f"{pad('', 18)}码字长度 {L}；−log2 p = {[round(-math.log2(q), 2) for q in p]}") print("\n观察 Observations：") print(" 1. 香农（1948）：概率为 p 的符号，最优码长是 −log2 p 个 bit；平均下来就是熵 H。") print(" 这把'预测'和'压缩'连在了一起：模型给真实符号的概率越高，编码它花的 bit 越少。") print(" 2. 概率都是 2 的幂时，哈夫曼编码正好达到熵。分布很偏时（0.9 的那个符号只值 0.15 bit），") print(" 哈夫曼编码每个符号仍要至少 1 bit，平均 1.15 bit，比熵多一倍多。") print(" 3. 算术编码把整条消息编成一个数，不受'每个符号整数个 bit'的限制，实测码长几乎等于熵。") print(" 这就是为什么'交叉熵 = 码长'在实践中是字面意义上成立的。") # ============================================================================= # 演示 2：压缩验证集 # ============================================================================= def demo2_compress(world, tok, base): print_header("演示 2：压缩第 1 讲的验证集文本（每个字符原本占 1 字节 = 8 bit）") val = world.val_text() raw = len(val.encode()) print(f"验证集：{len(val):,} 个字符，{raw:,} 字节") print(f"\n{pad('压缩器', 30)}{pad('理论 bit/字符', 14, '>')}{pad('实际 bit/字符', 14, '>')}{pad('压缩后字节', 12, '>')}" f"{pad('压缩率', 9, '>')}{pad('解压一致', 10, '>')}") rows = [] uniform = [[1 / len(ALPHABET)] * len(ALPHABET)] * len(val) rows.append(("均匀分布（44 个字符）", uniform)) train = world.train_text() for n in (1, 3, 5): lm = CharNGramLM(n, ALPHABET, k=0.01).fit(train) rows.append((f"{n}-gram（在训练集上数数）", ngram_distributions(lm, val))) rows.append(("base model（TinyGPT，109k 参数）", gpt_distributions(base, tok, val))) for name, dists in rows: bits = encode_text(val, ALPHABET, dists) ok = decode_text(bits, ALPHABET, dists) == val ideal = ideal_bits(val, ALPHABET, dists) / len(val) nbytes = math.ceil(len(bits) / 8) print(f"{pad(name, 30)}{pad(f'{ideal:.4f}', 14, '>')}{pad(f'{len(bits) / len(val):.4f}', 14, '>')}" f"{pad(f'{nbytes:,}', 12, '>')}{pad(f'{raw / nbytes:.1f}×', 9, '>')}{pad('是' if ok else '否', 10, '>')}") data = val.encode() for name, fn in [("gzip -9", lambda b: gzip.compress(b, 9, mtime=0)), ("bz2 -9", lambda b: bz2.compress(b, 9)), ("lzma（xz）-9", lambda b: lzma.compress(b, preset=9))]: c = fn(data) print(f"{pad(name, 30)}{pad('—', 14, '>')}{pad(f'{8 * len(c) / len(val):.4f}', 14, '>')}{pad(f'{len(c):,}', 12, '>')}" f"{pad(f'{raw / len(c):.1f}×', 9, '>')}{pad('是', 10, '>')}") print("\n观察 Observations：") print(" 1. 实际码长和理论码长（交叉熵）只差千分之几 bit/字符：模型的交叉熵**就是**它作为压缩器的成绩。") print(" 2. 模型越好，文件越小：均匀分布 5.46 bit/字符（比原始的 8 bit 好，因为只有 44 个字符），") print(" 3-gram、5-gram 依次变好，base model 只要 0.36 bit/字符——把验证集压成原来的约 1/22。") print(" 3. 通用压缩器 gzip / bz2 / lzma 只能找到 1.2–1.5 bit/字符：它们不'懂'这个世界的规律") print(" （谁住在哪里、评论的评分规则），只能利用重复的片段。base model 比 bz2 还小 3 倍多。") print(" 4. 注意 base model 的 0.36 比第 1 讲报告的验证集 0.409 bit/字符更好：这里每个字符都有至少 48 个字符的上下文，") print(" 而 bits_per_char 把文本切成互不重叠的窗口，每个窗口开头的字符几乎没有上下文。") # ============================================================================= # 演示 3：边解边预测 # ============================================================================= def next_distribution(model, ids): """给定已经解出的字符编号 ids（含开头的换行），返回下一个字符的概率分布（只看最后 block_size 个）。""" x = torch.tensor([ids[-model.cfg.block_size:]], dtype=torch.long) with torch.no_grad(): return torch.softmax(model(x)[0, -1].double(), dim=-1).tolist() def demo3_incremental(world, tok, base): print_header("演示 3：解码器事先不知道文本——每解出一个字符，才算下一个位置的预测") text = world.val_text()[:600] nl = tok.stoi["\n"] enc = ArithmeticEncoder() ids = [nl] for ch in text: # 编码器：一个字符一个字符地预测、编码 enc.encode(quantize(next_distribution(base, ids)), tok.stoi[ch]) ids.append(tok.stoi[ch]) bits = enc.finish() dec = ArithmeticDecoder(bits) out = [nl] for _ in range(len(text)): # 解码器：只用自己已经解出来的字符做预测 out.append(dec.decode(quantize(next_distribution(base, out)))) decoded = tok.decode(out[1:]) print(f"前 {len(text)} 个字符：压成 {len(bits)} bit（{len(bits) / len(text):.3f} bit/字符）；解码器边解边预测，结果与原文一致：{decoded == text}") print(f"解出的开头：{decoded[:70]!r}") print("\n观察 Observations：") print(" 1. 解码器手里只有比特流和同一个模型。它从开头的换行出发，预测第一个字符的分布、解出第一个字符，") print(" 再用'换行 + 第一个字符'预测第二个……编码器和解码器在每一步看到完全相同的上下文，") print(" 算出完全相同的分布，所以解码是精确的。") print(" 2. 这也意味着：**发送方和接收方必须有同一个模型**。如果接收方没有模型，就得把模型也发过去——演示 4。") # ============================================================================= # 演示 4：模型的代价 # ============================================================================= def demo4_model_cost(world, tok, base): print_header("演示 4：用 base model 压缩它自己的训练文本——别忘了模型本身也要发过去") train = world.train_text() dists = gpt_distributions(base, tok, train) bits = ideal_bits(train, ALPHABET, dists) n = count_parameters(base) raw = len(train.encode()) print(f"训练文本：{len(train):,} 个字符（{raw / 1024:.1f} KB）") print(f" base model 的码长：{bits / len(train):.4f} bit/字符 → {bits / 8 / 1024:.1f} KB") for b in (32, 16, 8): model_kb = n * b / 8 / 1024 print(f" 加上模型参数（{n:,} 个 × {b} bit = {model_kb:.1f} KB）→ 合计 {(bits / 8 / 1024) + model_kb:.1f} KB") for name, c in [("gzip -9", gzip.compress(train.encode(), 9, mtime=0)), ("bz2 -9", bz2.compress(train.encode(), 9)), ("lzma -9", lzma.compress(train.encode(), preset=9))]: print(f" 对照：{pad(name, 8)}{len(c) / 1024:.1f} KB") print("\n观察 Observations：") print(" 1. base model 把训练文本压到 0.350 bit/字符，和验证集（0.356）几乎一样：这个世界的文本高度规则，") print(" 模型学到的是规律，而不是逐字背下训练集——所以没见过的验证集也压得一样好。") print(" 2. 但它本身有 10.9 万个参数：按 32 位存是 426 KB，比原始文本（257 KB）还大，比 bz2 的结果大二十多倍。") print(" 这种'先发模型、再发用模型压缩的数据'的两段式编码（two-part code），对小数据集是亏的。") print(" 3. 真实的大模型也一样：一个 65B 参数的模型存下来要一百多 GB。只要把模型当成要发送的东西，") print(" '模型越大压缩越好'就不再显然——代码 02 介绍一种不需要发送模型的编码方式。") if __name__ == "__main__": world, tok = build_world(), CharTokenizer(ALPHABET) base, _ = load_or_train_base_model(world, tok) demo1_codes() demo2_compress(world, tok, base) demo3_incremental(world, tok, base) demo4_model_cost(world, tok, base)