"""
=============================================================================
前置讲 P1 · 代码 03：token 怎样影响提示 —— 语言、数字、空格、字母
=============================================================================

对应讲义：lectures/prereq1_text_bytes_tokens.md §4

【这个脚本要回答的问题】
    1. 同样的意思，用中文写和用英文写，token 数差多少？差距从哪里来？（演示 1："token 税"）
    2. 数字是怎样被切开的？为什么这会让模型算错数？（演示 2）
    3. 提示末尾多一个空格，为什么可能让结果变差？（演示 3）
    4. 为什么模型数不清一个单词里有几个字母？（演示 4）

【怎么读输出】
    这里训练了几个小分词器（各 1600 次合并），数字只用来说明**机制**；真实模型的分词器大得多，
    具体比例不同，但方向一致（讲义里给出了论文里真实分词器的数字）。

【运行】python3 03_token_effects.py     （约 25 秒：要训练 5 个分词器）
【依赖】只用 Python 标准库；import 同目录的 text_lab.py 和第 1 讲的 console.py。
"""

import os
import sys
from collections import Counter

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
from console import pct, print_header, print_observations, print_table  # noqa: E402
from text_lab import ByteBPE, _lexicon, bilingual_corpus  # noqa: E402

# 12 对意思相同的中英文句子（自己写的，不来自任何外部文本）
PARALLEL = [
    ("The soup was cold, but the bread was fresh.", "汤凉了，但是面包很新鲜。"),
    ("What is the capital of this country?", "这个国家的首都是哪里？"),
    ("Please answer with a single word.", "请只用一个词回答。"),
    ("The staff were friendly and the room was clean.", "服务员很热情，房间也很干净。"),
    ("Think step by step before you answer.", "回答之前请一步一步地思考。"),
    ("Summarize the following article in three sentences.", "请用三句话总结下面这篇文章。"),
    ("The price is too high for such a small room.", "这么小的房间，价格太贵了。"),
    ("If you do not know the answer, say that you do not know.", "如果你不知道答案，就说你不知道。"),
    ("The king lives in the palace and wears a crown.", "国王住在宫殿里，戴着王冠。"),
    ("Translate the next sentence into English.", "把下一句话翻译成英文。"),
    ("The meeting was moved to next Tuesday afternoon.", "会议改到了下周二下午。"),
    ("Return the result as a JSON object.", "请把结果写成一个 JSON 对象返回。"),
]


def train(en_share, digits="group3", n_merges=1600):
    """在指定中英文比例的语料上训练一个分词器。"""
    en, zh = bilingual_corpus(en_share=en_share)
    return ByteBPE(digits=digits).train(en + zh, n_merges=n_merges)


# ---------------------------------------------------------------------------
def demo1_language_tax(tokenizers):
    """演示 1：同样的意思，不同语言的 token 数。"""
    print_header("演示 1：'token 税'——同样的意思，换一种语言要多少 token？")
    rows, ratio = [], {}
    for name, bpe in tokenizers.items():
        en = sum(len(bpe.encode(e)) for e, _ in PARALLEL)
        zh = sum(len(bpe.encode(z)) for _, z in PARALLEL)
        en_c = sum(len(e) for e, _ in PARALLEL)
        zh_c = sum(len(z) for _, z in PARALLEL)
        ratio[name] = zh / en
        rows.append([name, en, zh, f"{zh / en:.2f}", f"{en_c / en:.2f}", f"{zh_c / zh:.2f}"])
    print(f"{len(PARALLEL)} 对意思相同的句子（英文共 {sum(len(e) for e, _ in PARALLEL)} 个字符，"
          f"中文共 {sum(len(z) for _, z in PARALLEL)} 个字符）")
    print_table(["分词器的训练语料", "英文 token", "中文 token", "中文 / 英文", "英文：字符 / token", "中文：字符 / token"], rows)
    bpe = tokenizers["英文 90% / 中文 10%"]
    e, z = PARALLEL[6]
    print(f"\n  例子（英文 90% 的分词器）：\n    {bpe.pieces(e)}\n    {bpe.pieces(z)}")
    names = list(tokenizers)
    en_tok = {n: sum(len(b.encode(e)) for e, _ in PARALLEL) for n, b in tokenizers.items()}
    print_observations([
        f"随着训练语料里中文的比例从 10% 升到 90%，中文 / 英文的 token 比从 {ratio[names[0]]:.2f} 降到 {ratio[names[2]]:.2f}，"
        f"英文 token 从 {en_tok[names[0]]} 升到 {en_tok[names[2]]}：哪种语言在分词器的训练语料里多，哪种语言就'便宜'。"
        "token 的多少不是语言本身的属性，而是**分词器训练语料的配比**和**词表大小**共同决定的。",
        f"在这个只有 1856 个 token 的小分词器上，中文并没有比英文贵：同样的意思中文只用了 "
        f"{sum(len(z) for _, z in PARALLEL)} 个字符（英文 {sum(len(e) for e, _ in PARALLEL)} 个），常用字又大多学成了"
        "单个 token；而英文单词还常被切成好几段（每个 token 平均只有 2.3 个字符左右）。",
        "真实分词器的词表大几十倍：按 OpenAI 文档的经验值，英文大约每 4 个字符一个 token，常用词几乎都是一个 token；"
        "如果中文每个字仍要一个甚至更多 token，同样的意思写成中文就会更贵。Petrov et al.（2023）在真实分词器上测到，"
        "同样内容的 token 数在不同语言之间最多能差 15 倍，受影响最大的是训练语料里少见的语言。",
        "【提示工程联系】API 按 token 计费，上下文窗口按 token 计算：同样的任务换一种语言，费用、延迟和"
        "'能放进多少背景资料'都会变。要知道一段提示到底多少 token，用你所用模型的分词器（或 API 的计数接口）去数，不要估。",
    ])


# ---------------------------------------------------------------------------
def split_pattern(bpe, s):
    """一个数字被切成的各段的长度，例如 '123456' → '3+3'。"""
    return "+".join(str(len(bpe.token_str(i))) for i in bpe.encode(s))


def demo2_numbers(digit_tokenizers):
    """演示 2：数字的切法。"""
    print_header("演示 2：数字是怎样被切开的？")
    samples = ["7", "42", "2026", "1024", "123456", "3.14159", "1,000,000", "1000000"]
    rows = []
    for s in samples:
        rows.append([s] + [" | ".join(b.pieces(s)) for b in digit_tokenizers.values()])
    print_table(["数字"] + list(digit_tokenizers), rows, aligns=["<"] * 4)
    print("\n所有 4 位数（1000–9999）的切法有几种？各占多少？")
    rows = []
    for name, b in digit_tokenizers.items():
        cnt = Counter(split_pattern(b, str(n)) for n in range(1000, 10000))
        top = ", ".join(f"{p}: {pct(c, 9000)}" for p, c in cnt.most_common(4))
        rows.append([name, len(cnt), top])
    print_table(["预分词规则", "切法种类", "最常见的切法（各段的位数）"], rows, aligns=["<", ">", "<"])
    print_observations([
        "'每位数字单独一个 token'最规整：任何数字的切法都一样，代价是序列变长。",
        "'最多 3 位一块'的预分词是**从左往右**分组的：'2026' 先被切成 '202' 和 '6' 两块（而不是按千位的 '2' 和 '026'），"
        "BPE 再在块内合并；'1,000,000' 和 '1000000' 的切法也完全不同。token 的边界和数位（个十百千）对不齐。",
        "'让 BPE 自己学'得到的切法最不规则：训练语料里常见的年份、价格（'2026'、'1024'）可能是一个 token，"
        "相邻的数字却被切成别的样子。模型做加法时，同一个数位在不同的数里落在 token 的不同位置上，"
        "这是早期 LLM 算术不稳的原因之一（Singh & Strouse 2024 发现，GPT-3.5/4 上把数字写成"
        "从右往左每 3 位一组——也就是加上千位分隔符——能明显提高算术准确率）。",
        "【提示工程联系】要模型做精确计算时，与其指望它'心算'，不如让它调用计算工具（第 8 讲）；"
        "实在要它算，统一数字的写法（例如都加千位分隔符）比随手写更稳。",
    ])


# ---------------------------------------------------------------------------
def demo3_spaces(bpe, train_texts):
    """演示 3：空格属于后面的那个词。"""
    print_header("演示 3：空格属于后面那个词——提示末尾多一个空格会怎样")
    for s in ["capital", " capital", "The capital of volel is grorken.", "The capital of volel is ", "The capital of volel is"]:
        print(f"  {s!r:38s} → {bpe.pieces(s)}")
    # 训练语料里，一个单独的 ' ' token 后面紧跟着字母开头的 token，有多少次？
    after_space = Counter()
    word_tokens = 0
    for t in train_texts:
        ids = bpe.encode(t)
        for a, b in zip(ids, ids[1:]):
            sa, sb = bpe.token_str(a), bpe.token_str(b)
            if sb[:1].isalpha():
                word_tokens += 1
                if sa == " ":
                    after_space["单独的空格 token 后面跟字母"] += 1
            if sb.startswith(" ") and sb[1:2].isalpha():
                after_space["带前导空格的词 token"] += 1
    print(f"\n  训练语料里：带前导空格的词 token 出现 {after_space['带前导空格的词 token']:,} 次；"
          f"'单独一个空格 token，后面紧跟字母开头的 token' 出现 {after_space['单独的空格 token 后面跟字母']} 次")
    print_observations([
        "预分词把空格粘在后面那个词的**前面**：' capital' 是一个 token，'capital' 是另一个。",
        "提示以 'is ' 结尾（多了一个空格）时，最后一个 token 变成了单独的 ' '。在训练语料里，单独的空格 token 后面"
        "几乎从来不会紧跟一个词——模型这时要写的 'grorken'（不带空格的版本）是它很少见过的续写方式。",
        "这就是为什么提示末尾的空格、换行会悄悄影响结果。有些推理库会做'token 修复'（token healing）："
        "把提示最后一个 token 退回去，让模型自己重新生成它。最简单的做法是：**不要让提示以空格结尾**。",
    ])


# ---------------------------------------------------------------------------
def demo4_letters(bpe):
    """演示 4：模型看到的是 token，不是字母。"""
    print_header("演示 4：模型看到的是 token，不是字母——'有几个 r？'")
    en_lex, _ = _lexicon()
    frequent, rare = en_lex[0], en_lex[2500]   # 长尾词典里最常见的词，和排名第 2501 的罕见词
    rows = []
    for w in ["strawberry", "capital", "Mississippi", frequent, rare, "提示工程"]:
        ids = bpe.encode(w)
        rows.append([w, " | ".join(bpe.pieces(w)), len(ids), ", ".join(str(i) for i in ids)])
    print_table(["单词", "token 片段", "token 数", "模型实际看到的 id"], rows, aligns=["<", "<", ">", "<"])
    rev = "yrrebwarts"
    print(f"\n  'strawberry' 反过来写是 {rev!r}：{bpe.pieces('strawberry')} → {bpe.pieces(rev)}")
    print_observations([
        "模型看到的是右边那一列整数。'strawberry' 里有几个 r？对模型来说，这要求它记得每个 token id 各由哪些字母组成，"
        "再把它们加起来——这是一个它从来没被直接训练过的任务。这就是 LLM 数不清字母、拼写和倒写单词容易出错的根源。",
        f"长尾词典里最常见的假词 {frequent!r} 只有 {len(bpe.encode(frequent))} 个 token，排名第 2501 的罕见假词 {rare!r} "
        f"被切成了 {len(bpe.encode(rare))} 段：越罕见的词切得越碎，模型对它的'拼写'也越没把握。",
        "把一个词倒过来写，得到的 token 和原来的毫无关系：倒写、按字母排序、数字母这类'字符级'任务，"
        "在 token 层面都很难。",
        "【提示工程联系】需要字符级操作时，让模型先把词**拆成一个字母一个字母**（'s-t-r-a-w-…'）再处理，"
        "或者干脆交给代码去做（第 8 讲的工具使用、第 11 讲的代码生成）。",
    ])


def main():
    print("训练 3 个不同语料配比的分词器 + 3 种数字规则的分词器（共 5 个，各 1600 次合并）……")
    tokenizers = {
        "英文 90% / 中文 10%": train(0.9),
        "英文 50% / 中文 50%": train(0.5),
        "英文 10% / 中文 90%": train(0.1),
    }
    digit_tokenizers = {
        "每位数字单独 single": train(0.9, digits="single"),
        "最多 3 位一块 group3": tokenizers["英文 90% / 中文 10%"],
        "让 BPE 自己学 any": train(0.9, digits="any"),
    }
    demo1_language_tax(tokenizers)
    demo2_numbers(digit_tokenizers)
    en, zh = bilingual_corpus(0.9)
    demo3_spaces(tokenizers["英文 90% / 中文 10%"], en)
    demo4_letters(tokenizers["英文 90% / 中文 10%"])


if __name__ == "__main__":
    main()
