nanoGPT复现——prepare拆解（自己构建词表 VS tiktoken）

在nanoGPT的data文件夹有两个很相似的文件夹结构：shakespeare和shakespeare-char，这两种都是对shakespeare数据集的处理，但是shakespeare使用的是tiktoken对文字进行编码，另一个则是使用自己构建的词表

一、shakespeare-char（自己构建词表）

数据获取

data_path = os.path.join(os.path.dirname(__file__), 'input.txt')
if not os.path.exists(data_path):
    url = 'https://cdn.jsdelivr.net/gh/karpathy/char-rnn@master/data/tinyshakespeare/input.txt'
    with open(data_path, 'w', encoding='utf-8') as f:
        f.write(requests.get(url).text)
with open(data_path, 'r', encoding='utf-8') as f:
    data = f.read()

我这里在运行的时候是没有办法直接下载的，如果出现这个情况就直接打开网址手动下载就好

构建词表

chars = sorted(list(set(data)))
stoi = {s: i for i, s in enumerate(chars)}
itos = {i: s for i, s in enumerate(chars)}

def encode(x):
    return [stoi[s] for s in x]
def decode(l):
    return ''.join([itos[i] for i in l])

划分训练集和测试集

n = len(data)
train_data = data[: int(0.9 * n)]
val_data = data[int(0.9 *n):]
train_idx = encode(train_data)
val_idx = encode(val_data)

对训练集和测试集分别编码

train_idx = np.array(train_idx, dtype=np.uint16)
val_idx = np.array(val_idx, dtype=np.uint16)
train_idx.tofile(os.path.join(os.path.dirname(__file__), 'train.bin'))
val_idx.tofile(os.path.join(os.path.dirname(__file__), 'val.bin'))

保存词表为meta.pkl文件（在sample.py中会用）

meta = {
    'voavb_size': len(chars),
    'itos': itos,
    'stoi': stoi
}
with open(os.path.join(os.path.dirname(__file__), 'meta.pkl'), 'wb') as f:
    pickle.dump(meta, f)
print('finish')

二、shakespeare（利用tiktoken）

数据加载、划分数据集的部分都相同，就不再赘述了

数据编码

enc = tiktoken.get_encoding('gpt2')
train_ids = enc.encode_ordinary(train_data)
val_ids = enc.encode_ordinary(val_data)
print(f"train has {len(train_ids):,} tokens")
print(f"val has {len(val_ids):,} tokens")

保存数据

train_ids = np.array(train_ids, dtype=np.uint16)
val_ids = np.array(val_ids, dtype=np.uint16)
val_ids.tofile(os.path.join(os.path.dirname(__file__), 'val.bin'))
train_ids.tofile(os.path.join(os.path.dirname(__file__), 'train.bin'))

三、关于保存数据的几种方式对比

不知道大家发现没有，就这几十行代码中有三种文件读写方式

（1）f.write/f.read

直接读写字符串或字节流，不涉及格式解释（如txt）

（2）val_ids.tofile

原始二进制存储（如bin）,但不保存shape需要提前知道数据格式

（3）pickle.dump

把任意Python对象（列表、字典、类、模型等）序列化或二进制流