# %%ai openai-chat-custom:Qwen3.5-9B-Q4_K_M.gguf
# 用一句话解释: 为什么LLM需要把文本变成数字ID？

import re

class SimpleTokenizer:
    """按空格+标点分词的简单实现"""
    def __init__(self, text):
        # 构建词表
        words = sorted(set(re.findall(r"[\w']+|[.,!?;:]", text.lower())))
        words.insert(0, "<unk>")  # 未知词
        words.insert(1, "<eos>")  # 句子结束
        self.vocab = {w: i for i, w in enumerate(words)}
        self.id_to_token = {i: w for w, i in self.vocab.items()}
    
    def encode(self, text):
        words = re.findall(r"[\w']+|[.,!?;:]", text.lower())
        return [self.vocab.get(w, self.vocab["<unk>"]) for w in words]
    
    def decode(self, ids):
        return " ".join(self.id_to_token[i] for i in ids)

# 测试
text = "Hello, world! This is a tokenizer."
tokenizer = SimpleTokenizer(text)
ids = tokenizer.encode(text)
print(f"词表大小: {len(tokenizer.vocab)}")
print(f"编码: {text} → {ids}")
print(f"解码: {ids} → {tokenizer.decode(ids)}")

try:
    import tiktoken
    print("tiktoken 已安装")
except ImportError:
    print("安装 tiktoken...")
    import subprocess, sys
    subprocess.check_call([sys.executable, "-m", "pip", "install", "tiktoken"])
    import tiktoken

# GPT-2 使用的 BPE 分词器
enc = tiktoken.get_encoding("gpt2")

text = "Hello, world! This is a tokenizer in action."
ids = enc.encode(text)
tokens = [enc.decode_single_token_bytes(t) for t in ids]

print(f"原始文本: {text}")
print(f"BPE 词表大小: {enc.n_vocab}")
print(f"编码 IDs: {ids}")
print(f"token 字节: {[t.decode('utf-8', errors='replace') for t in tokens]}")
print(f"解码: {enc.decode(ids)}")

# 看看 BPE 如何处理罕见词
rare_words = ["Python", "transformer", "注意力机制", "GPT", "tokenization"]
for word in rare_words:
    ids = enc.encode(word)
    tokens = [enc.decode_single_token_bytes(t).decode('utf-8', errors='replace') for t in ids]
    print(f"'{word}' → {ids} → {' | '.join(tokens)}")

import torch

def create_dataloader(text, batch_size=2, max_length=6, stride=3, shuffle=True):
    """
    滑动窗口数据加载器
    - text: 输入文本
    - max_length: 每个样本的token数
    - stride: 窗口滑动步长 (stride < max_length 时有重叠)
    """
    enc = tiktoken.get_encoding("gpt2")
    token_ids = enc.encode(text)
    
    inputs, targets = [], []
    for i in range(0, len(token_ids) - max_length, stride):
        inputs.append(token_ids[i:i+max_length])
        targets.append(token_ids[i+1:i+max_length+1])  # 右移一位
    
    dataset = torch.utils.data.TensorDataset(
        torch.tensor(inputs),
        torch.tensor(targets)
    )
    return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)

# 测试
sample = "The quick brown fox jumps over the lazy dog. Machine learning is transforming the world."
loader = create_dataloader(sample, batch_size=2, max_length=8, stride=4)

for batch_idx, (x, y) in enumerate(loader):
    print(f"批次 {batch_idx}:")
    print(f"  input:  {x.tolist()}")
    print(f"  target: {y.tolist()}")
    if batch_idx >= 2: break

import torch.nn as nn

vocab_size = 50257  # GPT-2 词表
emb_dim = 256       # 嵌入维度 (越小越好演示)
context_len = 8     # 上下文长度

# 1. Token Embedding: 每个 token ID → 256维向量
token_embedding = nn.Embedding(vocab_size, emb_dim)

# 2. 位置 Embedding: 每个位置 → 256维向量 (可学习的)
pos_embedding = nn.Embedding(context_len, emb_dim)

# 前向计算
sample_ids = torch.tensor([[15496, 11, 995, 0, 345, 23, 567, 890]])  # (batch=1, seq=8)
token_embeds = token_embedding(sample_ids)     # (1, 8, 256)
pos_embeds = pos_embedding(torch.arange(context_len).unsqueeze(0))  # (1, 8, 256)
input_embeds = token_embeds + pos_embeds       # (1, 8, 256)

print(f"Token Embedding 形状: {token_embeds.shape}")
print(f"位置 Embedding 形状: {pos_embeds.shape}")
print(f"最终输入形状: {input_embeds.shape}")
print(f"\n第一个token的嵌入向量 (前10维): {input_embeds[0, 0, :10]}")

# %%ai openai-chat-custom:Qwen3.5-9B-Q4_K_M.gguf --format markdown
# 对比解释: 为什么GPT使用可学习的位置嵌入而不是Transformer原文的三角函数位置编码？各有什么优缺点？