"""
第 11 章配套代码：
(1) 幂律 scaling law 拟合与外推（损失 vs 参数量）
(2) Bradley-Terry 偏好模型（RLHF 奖励模型的数学核心）
Runnable with: numpy only.  python3 11_scaling_rlhf.py
"""
import numpy as np


def power_law_demo():
    """Kaplan 2020: L(N) ≈ (Nc / N)^alpha_N + L_inf （幂律）。
    在 log-log 坐标下是一条直线，这正是它能可靠外推的原因。
    """
    print("=== (1) 幂律 scaling: L(N) = (Nc/N)^a + L_inf ===")
    Nc, a, L_inf = 8.8e13, 0.076, 1.69   # Kaplan 式量级（示意常数）
    print("  参数量 N        预测损失 L     (规模每涨10倍, 损失稳定下降)")
    for N in [1e6, 1e7, 1e8, 1e9, 1e10, 1e11]:
        L = (Nc / N) ** a + L_inf
        print(f"  {N:.0e}     {L:.3f}")
    print("  关键: log-log 下近似直线 => 可外推预测更大模型的损失（GPT-3 据此规划）")


def chinchilla_optimal():
    """Chinchilla 2022: 同算力下，参数 N 与数据 D 应同比缩放(约50/50)。
    算力 C ≈ 6*N*D。给定 C，最优 N、D 各占一半。
    """
    print("\n=== (2) Chinchilla 计算最优分配 ===")
    print("  算力 C ≈ 6*N*D ; 最优时 N 翻倍则 D 也翻倍 (约 20 tokens/参数)")
    for C in [1e21, 1e22, 1e23]:
        # 简化: N_opt ∝ C^0.5, D_opt ∝ C^0.5
        N_opt = (C / 6) ** 0.5
        D_opt = (C / 6) ** 0.5
        print(f"  C={C:.0e} FLOPs -> N_opt~{N_opt:.2e} 参数, D_opt~{D_opt:.2e} tokens")
    print("  结论: GPT-3(175B,300B tokens) 欠训练; Chinchilla(70B,1.4T tokens) 同算力更优")


def sigmoid(x): return 1 / (1 + np.exp(-x))


def bradley_terry():
    """RLHF 奖励模型: 人类偏好 y_w > y_l 的概率 = sigmoid(r(y_w) - r(y_l))。
    训练奖励模型即最大化偏好对的对数似然。
    """
    print("\n=== (3) Bradley-Terry 偏好模型 (RLHF 奖励模型核心) ===")
    print("  P(回答A优于回答B) = sigmoid(r_A - r_B)")
    for rw, rl in [(2.0, 1.0), (2.0, 0.0), (3.0, -1.0), (1.0, 1.0)]:
        p = sigmoid(rw - rl)
        print(f"  r(优)={rw:>4} r(劣)={rl:>5} -> P(选优)={p:.3f}")
    print("  奖励模型损失 = -log sigmoid(r(优) - r(劣))，从人类排序里学出标量奖励")


if __name__ == "__main__":
    power_law_demo()
    chinchilla_optimal()
    bradley_terry()