# -*- coding: utf-8 -*-
"""
tokenizer_play.py · 第 0002 课动手脚本
用真实的 GLM-4.5 词表（BPE, byte-level）切分文本，亲眼看看 token 长什么样。

用法：
    conda activate py310_qwenpaw
    python scripts/tokenizer_play.py                 # 切分内置示例
    python scripts/tokenizer_play.py "你的任意文本"    # 切分自己的文本

说明：
- 词表文件 tokenizer.json 从 ModelScope 下载（约 15MB），首次运行后缓存到本地
- 不需要 GPU、不需要联网账号；tokenizers 库你环境里已有（0.22.2）
"""

import sys
import urllib.request
from pathlib import Path

from tokenizers import Tokenizer

# Windows 编码兜底：无论控制台还是重定向，强制 UTF-8 输出
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")

# GLM-4.5 的 tokenizer.json（ModelScope 直链，已验证可达）
URL = "https://modelscope.cn/models/ZhipuAI/GLM-4.5/resolve/master/tokenizer.json"
CACHE = Path(__file__).parent / "_tokenizer_cache" / "glm45_tokenizer.json"


def load_tokenizer() -> Tokenizer:
    if CACHE.exists():
        print(f"[缓存命中] {CACHE}")
        return Tokenizer.from_file(str(CACHE))

    print("[首次运行] 从 ModelScope 下载 GLM-4.5 tokenizer.json ...")
    CACHE.parent.mkdir(parents=True, exist_ok=True)
    req = urllib.request.Request(URL, headers={"User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(req) as resp, open(CACHE, "wb") as f:
        f.write(resp.read())
    print(f"[下载完成] 缓存于 {CACHE}")
    return Tokenizer.from_file(str(CACHE))


def show(tok: Tokenizer, text: str) -> None:
    enc = tok.encode(text)
    print(f"\n原文: {text!r}")
    print(f"token 数: {len(enc.ids)}")
    print("-" * 60)
    for tid, tok_str in zip(enc.ids, enc.tokens):
        # byte-level BPE 的字节标记（Ġ = 空格，ą = 换行等）原样展示
        print(f"  id={tid:<7} {tok_str!r}")
    print("-" * 60)
    print(f"还原检查: {tok.decode(enc.ids, skip_special_tokens=False)!r}")


def main() -> None:
    tok = load_tokenizer()
    print(f"\n词表大小 vocab_size = {tok.get_vocab_size(True):,}")

    samples = [
        "strawberry",
        "大语言模型改变了世界",
        "3.1415926",
        "def quick_sort(arr):\n    if len(arr) <= 1:\n        return arr",
        "The unbelievable tokenizer retrokenized unbelievably.",
    ]
    if len(sys.argv) > 1:
        samples = [" ".join(sys.argv[1:])]

    for s in samples:
        show(tok, s)

    print("\n>>> 课堂任务：把你想测试的文本作为命令行参数传进来，")
    print('>>> 例如：python scripts/tokenizer_play.py "DeepSeek-V3 是 671B-A37B 的 MoE 模型"')


if __name__ == "__main__":
    main()