LLMs-from-scratch（dataloader）

代码链接：https://github.com/rasbt/LLMs-from-scratch/blob/main/ch02/01_main-chapter-code/dataloader.ipynb

《从零开始构建大型语言模型》一书的补充代码，作者：Sebastian Raschka 代码仓库：https://github.com/rasbt/LLMs-from-scratch

主要数据加载管道总结

完整的章节代码位于 ch02.ipynb。

这个笔记本包含了主要要点，即不包含中间步骤的数据加载管道。

本笔记本中使用的包：

# NBVAL_SKIP
from importlib.metadata import versionprint("torch version:", version("torch"))
print("tiktoken version:", version("tiktoken"))

torch version: 2.5.1+cu124
tiktoken version: 0.12.0

import tiktoken
import torch
from torch.utils.data import Dataset, DataLoaderclass GPTDatasetV1(Dataset):def __init__(self, txt, tokenizer, max_length, stride):self.input_ids = []self.target_ids = []# 对整个文本进行标记化token_ids = tokenizer.encode(txt, allowed_special={"<|endoftext|>"})# 使用滑动窗口将书籍分割成重叠的 max_length 长度序列for i in range(0, len(token_ids) - max_length, stride):input_chunk = token_ids[i:i + max_length]target_chunk = token_ids[i + 1: i + max_length + 1]self.input_ids.append(torch.tensor(input_chunk))self.target_ids.append(torch.tensor(target_chunk))def __len__(self):return len(self.input_ids)def __getitem__(self, idx):return self.input_ids[idx], self.target_ids[idx]def create_dataloader_v1(txt, batch_size, max_length, stride,shuffle=True, drop_last=True, num_workers=0):# 初始化标记器tokenizer = tiktoken.get_encoding("gpt2")# 创建数据集dataset = GPTDatasetV1(txt, tokenizer, max_length, stride)# 创建数据加载器dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers)return dataloader# 读取文本文件
with open("the-verdict.txt", "r", encoding="utf-8") as f:raw_text = f.read()# 定义模型参数
vocab_size = 50257      # 词汇表大小
output_dim = 256        # 输出维度
context_length = 1024   # 上下文长度# 创建嵌入层
token_embedding_layer = torch.nn.Embedding(vocab_size, output_dim)  # 标记嵌入层
pos_embedding_layer = torch.nn.Embedding(context_length, output_dim)  # 位置嵌入层# 设置数据加载器参数
batch_size = 8
max_length = 4
dataloader = create_dataloader_v1(raw_text,batch_size=batch_size,max_length=max_length,stride=max_length
)

# 处理一个批次的数据
for batch in dataloader:x, y = batch  # x: 输入序列, y: 目标序列# 获取标记嵌入token_embeddings = token_embedding_layer(x)# 获取位置嵌入pos_embeddings = pos_embedding_layer(torch.arange(max_length))# 将标记嵌入和位置嵌入相加得到输入嵌入input_embeddings = token_embeddings + pos_embeddingsbreak  # 只处理第一个批次作为示例

print(input_embeddings.shape)  # 输出嵌入张量的形状

torch.Size([8, 4, 256])