手撸 chatgpt 大模型:滑动窗口原理和代码实现

大语言模型（LLM）的主要任务是：给定输入文本，它应该能够正确预测下一个单词。例如，给定句子 “llm is used to generate text”，LLM 的输入和输出如下：

llm -> is

llm is -> used

llm is used -> generate

llm is used generate -> text

可以看出，这类似于使用滑动窗口将输入文本包含其中，期望的输出是窗口右边界之外的下一个单词。需要注意的是，在此之前，我们需要对文本进行分词（tokenize），也就是说，LLM 将接收到一个数字数组，并尝试预测下一个数字。

让我们通过代码实现滑动窗口策略。在上一节中，我们通过给定的 URL 下载了训练文本，结果是一个网页的 HTML 代码。我们需要跳过返回文本的很大一部分，直到找到我们需要用作训练文本的主要内容。如下图所示，我们需要跳过开头的一部分文本，才能到达主要内容：

然后使用以下代码，我们从选定文本中提取内容，并将其转换为 tokens：

import tiktoken
tokenizer = tiktoken.get_encoding('gpt2')

with open("fire-tongue.txt", "r", encoding="utf-8") as f:
  raw_text = f.read()

pos = raw_text.index("His investigation of the case of the man")
training_text = raw_text[pos:] 
encoded_text = tokenizer.encode(training_text)
print(len(encoded_text))

运行以上代码后，输出为 14429，表示原始文本被分词为 14429 个 token。

我们将窗口长度固定为 4，并生成输入文本和期望输出（即下一个单词）如下：

window_size = 4
# 窗口大小决定了有多少个 token 作为输入
x = encoded_text[:window_size]
# 右移一位以获得预测单词
y = encoded_text[1 : window_size+1]
print(f"x:  {x}")
print(f"y:       {y}")

for i in range(1, window_size + 1):
  input = encoded_text[:i]
  expect = encoded_text[i]
  print(input, "----->", expect)

for i in range(1, window_size + 1):
  input = encoded_text[:i]
  expect = encoded_text[i]
  print(tokenizer.decode(input), "----->" ,tokenizer.decode([expect]))

运行以上代码，我们得到以下输出：

x:  [6653, 3645, 286, 262]
y:       [3645, 286, 262, 1339]

[6653] -----> 3645
[6653, 3645] -----> 286
[6653, 3645, 286] -----> 262
[6653, 3645, 286, 262] -----> 1339

His ----->  investigation
His investigation ----->  of
His investigation of ----->  the
His investigation of the ----->  case

如我们所见，输入的窗口长度每次增加一个，给定窗口输入的期望输出是窗口右边界外的下一个单词。手动完成上述所有工作显然不现实，我们可以将任务交给现有的库，从而节省时间和精力。我们将使用的库是 torch，其 Dataset 和 DataLoader 可以帮助简化输入和输出的创建过程：

import torch
from torch.utils.data import Dataset, DataLoader

class GPTDatasetV1(Dataset):
  def __init__(self, input_text, tokenizer, window_size, shift):
    self.input_ids = []
    self.target_ids = []
    token_ids = tokenizer.encode(input_text)
    for i in range(0, len(token_ids) - window_size, shift):
      # 按 shift 的值向右移动窗口
      input_chunk = token_ids[i : i + window_size]
      target_chunk = token_ids[i+1 : i + window_size + 1]
      # tensor 本质上是一个向量
      self.input_ids.append(torch.tensor(input_chunk))
      self.target_ids.append(torch.tensor(target_chunk))

  def __len__(self):
    # 使 len() 可用于获取长度
    return len(self.input_ids)

  def __getitem__(self, idx):
    # 使 [] 可用于像数组一样获取元素
    return self.input_ids[idx], self.target_ids[idx]

def create_dataloader_v1(text, batch_size = 4, window_size = 256, shift = 128, shuffle = True, drop_last = True, num_workers = 0):
  tokenizer = tiktoken.get_encoding("gpt2")
  dataset = GPTDatasetV1(text, tokenizer, window_size, shift)
  """
  drop_last: 是否丢弃最后一个 batch（当 batch 的样本数不足 batch_size 时）
  num_workers: 用于运行 DataLoader 的线程数
  """
  dataloader = DataLoader(dataset, batch_size = batch_size, shuffle = shuffle,
                          drop_last = drop_last, num_workers = num_workers)
  return dataloader

with open("fire-tongue.txt", "r", encoding="utf-8") as f:
  raw_text = f.read()

dataloader = create_dataloader_v1(raw_text, batch_size = 1, window_size = 4, 
                                  shift = 1, shuffle = False)
data_iter = iter(dataloader)
first_batch = next(data_iter)
print(first_batch)

运行以上代码，我们得到以下输出：

[tensor([[   27,     0, 18227,  4177]]), tensor([[    0, 18227,  4177,    56]])]

可以看出，目标 ID 是在输入 ID 的基础上右移一位。当然，每次只发送一个输入和期望输出对的效率较低。在深度学习训练中，我们通常会收集一个包含多个样本的批次（batch）并同时发送给模型训练，这可以提高训练效率。我们可以通过以下方式增加批量大小：

# shift 设置为 4 表示给定输入后，模型需要预测接下来的四个单词
dataloader = create_dataloader_v1(raw_text, batch_size = 16, window_size=4, shift=4, shuffle = False)
data_iter = iter(dataloader)
inputs, targets = next(data_iter)
print(f"inputs\n: {inputs}")
print(f"outputs:\n: {targets}")

输出如下：

inputs
: tensor([[   27,     0, 18227,  4177],
        [   56, 11401, 27711,    29],
        [  198,    27,  6494,  1398],
        [ 2625, 16366,    12,  3919],
        [ 8457,     1, 42392,  2625],
        [  268,     1, 26672,  2625],
        [   75,  2213,  5320,   198],
        [   27,  2256,    29,   198],
        [   27, 28961, 34534,   316],
        [ 2625, 48504,    12,    23],
        [ 5320,   198,    27,  7839],
        [   29, 13543,    12,    51],
        [  506,   518,    14, 14126],
        [  352,   532, 11145,   271],
        [ 1668,    11,   262,  1479],
        [ 2691,  5888,  3556,  7839]])
outputs:
: tensor([[    0, 18227,  4177,    56],
        [11401, 27711,    29,   198],
        [   27,  6494,  1398,  2625],
        [16366,    12,  3919,  8457],
        [    1, 42392,  2625,   268],
        [    1, 26672,  2625,    75],
        [ 2213,  5320,   198,    27],
        [ 2256,    29,   198,    27],
        [28961, 34534,   316,  2625],
        [48504,    12,    23,  5320],
        [  198,    27,  7839,    29],
        [13543,    12,    51,   506],
        [  518,    14, 14126,   352],
        [  532, 11145,   271,  1668],
        [   11,   262,  1479,  2691],
        [ 5888,  3556,  7839,    29]])

更多经常内容请看这里