Transformer 革命：理解位元組級BPE演算法最佳化程式碼

print(tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str(python_code))

[('def', (0, 3)), ('Ġsay', (3, 7)), ('_', (7, 8)), ('hello', (8, 13)), ('():', (13, 16)), ('ĊĠĠĠ', (16, 20)), ('Ġprint', (20, 26)), ('("', (26, 28)), ('Hello', (28, 33)), (',', (33, 34)), ('ĠWorld', (34, 40)), ('!")', (40, 43)), ('Ġ#', (43, 45)), ('ĠPrint', (45, 51)), ('Ġit', (51, 54)), ('Ċ', (54, 55)), ('Ċ', (55, 56)), ('say', (56, 59)), ('_', (59, 60)), ('hello', (60, 65)), ('()', (65, 67)), ('Ċ', (67, 68))]

print(f"詞彙表大小: {len(tokenizer)}")
# 輸出: 詞彙表大小: 50257

print(tokenizer(python_code).tokens())

['def', 'Ġsay', '_', 'hello', '():', 'Ċ', 'Ġ', 'Ġ', 'Ġ', 'Ġprint', '("', 'Hello', ',', 'ĠWorld', '!"', ')', 'Ġ#', 'ĠPrint', 'Ġit', 'Ċ', 'Ċ', 'say', '_', 'hello', '()', 'Ċ']

tokens = sorted(tokenizer.vocab.items(), key=lambda x: len(x[0]), reverse=True)
print([f'{tokenizer.convert_tokens_to_string(t)}' for t, _ in tokens[:8]])

tokens = sorted(tokenizer.vocab.items(), key=lambda x: x[1], reverse=True)
print([f'{tokenizer.convert_tokens_to_string(t)}' for t, _ in tokens[:12]])

['<|endoftext|>', ' gazed', ' informants', ' Collider', ' regress', 'ominated', ' amplification', 'Compar', '..."', ' (/', 'Commission', ' Hitman']

from tqdm.auto import tqdm
length = 100000
dataset_name = 'transformersbook/codeparrot-train'
dataset = load_dataset(dataset_name, split="train", streaming=True)
iter_dataset = iter(dataset)

def batch_iterator(batch_size=10):
    for _ in tqdm(range(0, length, batch_size)):
        yield [next(iter_dataset)['content'] for _ in range(batch_size)]

new_tokenizer = tokenizer.train_new_from_iterator(batch_iterator(),
                                                 vocab_size=12500,
                                                 initial_alphabet=base_vocab)

tokens = sorted(new_tokenizer.vocab.items(), key=lambda x: x[1], reverse=False)
print([f'{tokenizer.convert_tokens_to_string(t)}' for t, _ in tokens[257:280]])

[' ', ' ', ' ', ' ', 'se', 'in', ' ', 're', 'on', 'te', '\n
', '\n ', 'or', 'st', 'de', '\n ', 'th', 'le', ' =', 'lf', 'self', 'me', 'al']

print([f'{new_tokenizer.convert_tokens_to_string(t)}' for t,_ in tokens[-12:]])

[' capt', ' embedded', ' regarding', 'Bundle', '355', ' recv', ' dmp', ' vault', ' Mongo', ' possibly', 'implementation', 'Matches']

special_tokens = {
    "additional_special_tokens": [
        "<python>", "<javascript>", "<function>", "<class>", "<docstring>"
    ]
}
tokenizer.add_special_tokens(special_tokens)

from tokenizers import pre_tokenizers

# 自定義前處理器，保留程式碼縮排
custom_pre_tokenizer = pre_tokenizers.Sequence([
    pre_tokenizers.Whitespace(),
    pre_tokenizers.Digits(individual_digits=False)
])
tokenizer.backend_tokenizer.pre_tokenizer = custom_pre_tokenizer

def post_process_code_tokens(tokens):
    # 合併連續的縮排標記
    processed_tokens = []
    i = 0
    while i < len(tokens):
        if tokens[i] == 'Ġ' and i+1 < len(tokens) and tokens[i+1] == 'Ġ':
            # 計算連續空格數
            count = 0
            while i < len(tokens) and tokens[i] == 'Ġ':
                count += 1
                i += 1
            # 增加單一縮排標記
            processed_tokens.append(f'<indent{count}>')
        else:
            processed_tokens.append(tokens[i])
            i += 1
    return processed_tokens

model.resize_token_embeddings(len(new_tokenizer))

# 使用相似標記的嵌入來初始化新標記
def initialize_new_tokens(model, tokenizer, old_tokenizer):
    # 取得舊詞彙表中不存在的新標記
    new_tokens = [t for t in tokenizer.vocab if t not in old_tokenizer.vocab]
    
    # 為每個新標記找到最相似的舊標記
    for new_token in new_tokens:
        # 尋找字元上最相似的舊標記
        similar_token = find_most_similar_token(new_token, old_tokenizer.vocab)
        # 使用相似標記的嵌入初始化新標記
        new_idx = tokenizer.convert_tokens_to_ids(new_token)
        similar_idx = old_tokenizer.convert_tokens_to_ids(similar_token)
        model.embeddings.word_embeddings.weight.data[new_idx] = \
            model.embeddings.word_embeddings.weight.data[similar_idx].clone()

def tokenize_function(examples):
    result = new_tokenizer(
        examples["content"],
        truncation=True,
        max_length=512,
        padding="max_length",
        return_tensors="pt"
    )
    return result

tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=["content"]
)

def evaluate_tokenizer(tokenizer, text_samples):
    original_lengths = [len(sample) for sample in text_samples]
    tokenized_lengths = [len(tokenizer.encode(sample)) for sample in text_samples]
    
    compression_ratio = sum(original_lengths) / sum(tokenized_lengths)
    print(f"平均壓縮比: {compression_ratio:.2f}")
    print(f"平均原始長度: {sum(original_lengths)/len(original_lengths):.2f} 字元")
    print(f"平均標記長度: {sum(tokenized_lengths)/len(tokenized_lengths):.2f} 標記")
    
    return compression_ratio

print(new_tokenizer(python_code).tokens())

['def', 'Ġs', 'ay', '_', 'hello', '():', 'ĊĠĠĠ', 'Ġprint', '("', 'Hello', ',',
'ĠWor', 'ld', '!")', 'Ġ#', 'ĠPrint', 'Ġit', 'Ċ', 'Ċ', 's', 'ay', '_', 'hello',
'()', 'Ċ']

import keyword
print(f'There are in total {len(keyword.kwlist)} Python keywords.')
for keyw in keyword.kwlist:
    if keyw not in new_tokenizer.vocab:
        print(f'No, keyword `{keyw}` is not in the vocabulary')

There are in total 35 Python keywords.
No, keyword `await` is not in the vocabulary
No, keyword `finally` is not in the vocabulary
No, keyword `nonlocal` is not in the vocabulary

length = 200000
new_tokenizer_larger = tokenizer.train_new_from_iterator(
    batch_iterator(),
    vocab_size=32768,
    initial_alphabet=base_vocab
)

tokens = sorted(new_tokenizer_larger.vocab.items(), key=lambda x: x[1],
                reverse=False)
print([f'{tokenizer.convert_tokens_to_string(t)}' for t, _ in tokens[-12:]]);

['lineEdit', 'spik', ' BC', 'pective', 'OTA', 'theus', 'FLUSH', ' excutils',
'00000002', ' DIVISION', 'CursorPosition', ' InfoBar']

print(new_tokenizer_larger(python_code).tokens())

['def', 'Ġsay', '_', 'hello', '():', 'ĊĠĠĠ', 'Ġprint', '("', 'Hello', ',',
'ĠWorld', '!")', 'Ġ#', 'ĠPrint', 'Ġit', 'Ċ', 'Ċ', 'say', '_', 'hello', '()',
'Ċ']

for keyw in keyword.kwlist:
    if keyw not in new_tokenizer_larger.vocab:
        print(f'No, keyword `{keyw}` is not in the vocabulary')

No, keyword `nonlocal` is not in the vocabulary

model_ckpt = "codeparrot"
org = "transformersbook"
new_tokenizer_larger.push_to_hub(model_ckpt, organization=org)

reloaded_tokenizer = AutoTokenizer.from_pretrained(org + "/" + model_ckpt)
print(reloaded_tokenizer(python_code).tokens())

['def', 'Ġsay', '_', 'hello', '():', 'ĊĠĠĠ', 'Ġprint', '("', 'Hello', ',',
'ĠWorld', '!")', 'Ġ#', 'ĠPrint', 'Ġit', 'Ċ', 'Ċ', 'say', '_', 'hello', '()',
'Ċ']

new_tokenizer.push_to_hub(model_ckpt+ "-small-vocabulary", organization=org)

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
config = AutoConfig.from_pretrained("gpt2-xl", vocab_size=len(tokenizer))
model = AutoModelForCausalLM.from_config(config)

print(f'GPT-2 (xl) size: {model_size(model)/1000**2:.1f}M parameters')

GPT-2 (xl) size: 1529.6M parameters

model.save_pretrained("models/" + model_ckpt, push_to_hub=True, organization=org)

tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
config_small = AutoConfig.from_pretrained("gpt2", vocab_size=len(tokenizer))
model_small = AutoModelForCausalLM.from_config(config_small)
print(f'GPT-2 size: {model_size(model_small)/1000**2:.1f}M parameters')

GPT-2 size: 111.0M parameters

model_small.save_pretrained("models/" + model_ckpt + "-small", push_to_hub=True, organization=org)

examples, total_characters, total_tokens = 500, 0, 0
dataset = load_dataset('transformersbook/codeparrot-train', split='train', streaming=True)
for _, example in tqdm(zip(range(examples), iter(dataset)), total=examples):
    total_characters += len(example['content'])
    total_tokens += len(tokenizer(example['content']).tokens())
characters_per_token = total_characters / total_tokens
print(characters_per_token)

3.6233025034779565

import torch
from torch.utils.data import IterableDataset

class ConstantLengthDataset(IterableDataset):
    def __init__(self, tokenizer, dataset, seq_length=1024,
                 num_of_sequences=1024, chars_per_token=3.6):
        self.tokenizer = tokenizer
        self.concat_token_id = tokenizer.eos_token_id
        self.dataset = dataset
        self.seq_length = seq_length
        self.input_characters = seq_length * chars_per_token * num_of_sequences
        
    def __iter__(self):
        iterator = iter(self.dataset)
        more_examples = True
        while more_examples:
            buffer, buffer_len = [], 0
            while True:
                if buffer_len >= self.input_characters:
                    m=f"Buffer full: {buffer_len}>={self.input_characters:.0f}"
                    print(m)
                    break
                try:
                    m=f"Fill buffer: {buffer_len}<{self.input_characters:.0f}"
                    print(m)
                    buffer.append(next(iterator)["content"])
                    buffer_len += len(buffer[-1])
                except StopIteration:
                    iterator = iter(self.dataset)
                    
            all_token_ids = []
            tokenized_inputs = self.tokenizer(buffer, truncation=False)
            for tokenized_input in tokenized_inputs["input_ids"]:
                all_token_ids.extend(tokenized_input + [self.concat_token_id])
                
            for i in range(0, len(all_token_ids), self.seq_length):
                input_ids = all_token_ids[i : i + self.seq_length]
                if len(input_ids) == self.seq_length:
                    yield torch.tensor(input_ids)

shuffled_dataset = dataset.shuffle(buffer_size=100)
constant_length_dataset = ConstantLengthDataset(tokenizer, shuffled_dataset, num_of_sequences=10)
dataset_iterator = iter(constant_length_dataset)
lengths = [len(b) for _, b in zip(range(5), dataset_iterator)]
print(f"Lengths of the sequences: {lengths}")

Fill buffer: 0<36864
Fill buffer: 3311<36864
Fill buffer: 9590<36864
Fill buffer: 22177<36864
Fill buffer: 25530<36864
Fill buffer: 31098<36864
Fill buffer: 32232<36864
Fill buffer: 33867<36864
Buffer full: 41172>=36864
Lengths of the sequences: [1024, 1024, 1024, 1024, 1024]

# 注意我們在建立ConstantLengthDataset之前先打亂了原始資料集
train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=args.seed)

import torch
import torch.nn.functional as F
from datasets import load_dataset
from accelerate import Accelerator

# 原始程式碼
# device = 'cpu'
# model = torch.nn.Transformer().to(device)

# 使用Accelerate的程式碼
accelerator = Accelerator()
model = torch.nn.Transformer()
optimizer = torch.optim.Adam(model.parameters())
dataset = load_dataset('my_dataset')
data = torch.utils.data.DataLoader(dataset, shuffle=True)
model, optimizer, data = accelerator.prepare(model, optimizer, data)

model.train()
for epoch in range(10):
    for source, targets in data:
        # 原始程式碼
        # source = source.to(device)
        # targets = targets.to(device)
        
        optimizer.zero_grad()
        output = model(source)
        loss = F.cross_entropy(output, targets)
        
        # 原始程式碼
        # loss.backward()
        
        # 使用Accelerate的程式碼
        accelerator.backward(loss)
        optimizer.step()

from argparse import Namespace

# 註解的引數對應於小型模型的設定
config = {
    "train_batch_size": 2,  # 小模型為12
    "valid_batch_size": 2,  # 小模型為12
    "weight_decay": 0.1,
    "shuffle_buffer": 1000,
    "learning_rate": 2e-4,  # 小模型為5e-4
    "lr_scheduler_type": "cosine",
    "num_warmup_steps": 750,  # 小模型為2000
    "gradient_accumulation_steps": 16,  # 小模型為1
    "max_train_steps": 50000,  # 小模型為150000
    "max_eval_steps": -1,
    "seq_length": 1024,
    "seed": 1,
    "save_checkpoint_steps": 50000  # 小模型為15000
}
args = Namespace(**config)

from torch.utils.tensorboard import SummaryWriter
import logging
import wandb

def setup_logging(project_name):
    logger = logging.getLogger(__name__)
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S", 
        level=logging.INFO, 
        handlers=[
            logging.FileHandler(f"log/debug_{accelerator.process_index}.log"),
            logging.StreamHandler()
        ]
    )
    
    if accelerator.is_main_process:  # 我們只想設定一次日誌
        wandb.init(project=project_name, config=args)
        run_name = wandb.run.name
        tb_writer = SummaryWriter()
        tb_writer.add_hparams(vars(args), {'0': 0})
        logger.setLevel(logging.INFO)
        datasets.utils.logging.set_verbosity_debug()
        transformers.utils.logging.set_verbosity_info()
    else:
        tb_writer = None
        run_name = ''
        logger.setLevel(logging.ERROR)
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()
        
    return logger, tb_writer, run_name

def log_metrics(step, metrics):
    logger.info(f"Step {step}: {metrics}")
    if accelerator.is_main_process:
        wandb.log(metrics)
        [tb_writer.add_scalar(k, v, step) for k, v in metrics.items()]

from torch.utils.data.dataloader import DataLoader

def create_dataloaders(dataset_name):
    train_data = load_dataset(dataset_name+'-train', split="train", streaming=True)
    train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=args.seed)
    
    valid_data = load_dataset(dataset_name+'-valid', split="validation", streaming=True)
    
    train_dataset = ConstantLengthDataset(tokenizer, train_data, seq_length=args.seq_length)
    valid_dataset = ConstantLengthDataset(tokenizer, valid_data, seq_length=args.seq_length)
    
    train_dataloader = DataLoader(train_dataset, batch_size=args.train_batch_size)
    eval_dataloader = DataLoader(valid_dataset, batch_size=args.valid_batch_size)
    
    return train_dataloader, eval_dataloader

def get_grouped_params(model, no_decay=["bias", "LayerNorm.weight"]):
    params_with_wd, params_without_wd = [], []
    for n, p in model.named_parameters():
        if any(nd in n for nd in no_decay):
            params_without_wd.append(p)
        else:
            params_with_wd.append(p)
    return [
        {'params': params_with_wd, 'weight_decay': args.weight_decay},
        {'params': params_without_wd, 'weight_decay': 0.0}
    ]

def evaluate():
    model.eval()
    losses = []
    for step, batch in enumerate(eval_dataloader):
        with torch.no_grad():
            outputs = model(batch, labels=batch)
            loss = outputs.loss.repeat(args.valid_batch_size)
            losses.append(accelerator.gather(loss))
        if args.max_eval_steps > 0 and step >= args.max_eval_steps: 
            break
            
    loss = torch.mean(torch.cat(losses))
    try:
        perplexity = torch.exp(loss)
    except OverflowError:
        perplexity = torch.tensor(float("inf"))
        
    return loss.item(), perplexity.item()