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

class GPTDataset(Dataset):
    def __init__(self,raw_text,tokenizer,window_size,step):
        """ 
            传入原始文本，分词器，窗口大小，步长
        """
        self.inputs = []
        self.targets = []
        ids = tokenizer.encode(raw_text)
        # 这里实际截取长度为5
        for i in range(0,len(ids) - window_size,step):
            inputs_chunk = ids[i:i+window_size]
            output_chunk = ids[i+1:i+window_size+1]

            self.inputs.append(inputs_chunk)
            self.targets.append(output_chunk)
            

    def __len__(self):
        if len(self.inputs) != len(self.targets):
            raise "输入输出shape不一致！"
        return len(self.inputs)
    
    def __getitem__(self, index):
        return self.inputs[index],self.targets[index]