model/dataset.py

from typing import Union

from torch.utils.data import Dataset
from torch import LongTensor, cuda
from transformers import PreTrainedTokenizerFast
from fastparquet import ParquetFile
from torch.utils.data import DataLoader
from datasets import load_dataset
import datasets
import pyarrow.parquet as pq
from numpy import array, int64
from numpy.random import shuffle

# import sys 
# sys.path.extend(['.', '..'])

from config import PROJECT_ROOT

class MyDataset(Dataset):

    def __init__(self, 
                parquet_file: str,
                tokenizer_dir: str,
                keep_in_memory: bool=False,
                max_seq_len: int=512,
                buffer_size: int=40960,
            ) -> None:
        '''
        keep_in_memory: 是否将parquet文件转换为pandas.DataFrame格式存放到内存, 
            False将使用迭代生成器(迭代生成器不支持打乱数据)，减少大数据集内存占用
        '''
        super().__init__()

        if cuda.device_count() >= 2 and not keep_in_memory:
            raise ValueError(f'多GPU时使用MyDataset，参数keep_in_memory必须=True，否则无法进行分布式训练. 当前keep_in_memory={keep_in_memory}')

        self.keep_in_memory = keep_in_memory
        self.max_seq_len = max_seq_len

        # 使用pyarrow.parquet读取，to_pandas、for遍历速度更快
        parquet_table = pq.read_table(parquet_file)

        # 获取数据集长度
        self.length = parquet_table.num_rows

        # 缓冲区大小不能超过数据长度
        self.buffer_size = self.length if buffer_size > self.length else buffer_size

        if keep_in_memory:
            # 转化为pandas放到内存中
            self.data = parquet_table.to_pandas()  
        else:
            self.data = parquet_table

        # 初始化tokenizer
        self.tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_dir)

        # 在这里初始化generator
        self.sample_generator = self.item_generator()
    
    def item_generator(self,) -> tuple:
        '''
        一条数据的生成器，防止大数据集OOM
        '''
                
        parquet_table = self.data

        # 生成器是死循环，不用退出，训练结束（epoch结束）会停止调用next()
        buffer_list = []
        while True:

            for prompt, response in zip(parquet_table['prompt'], parquet_table['response']):
                
                # 缓存数据不够，添加数据
                if len(buffer_list) < self.buffer_size:
                    buffer_list.append( (prompt.as_py(), response.as_py()) )
                    continue
                
                # 执行到这里，缓存区够了，打乱数据
                shuffle(buffer_list)
                for p, r in buffer_list:
                    # 在这里迭代
                    yield  p, r

                # 迭代完成，清空缓存区
                buffer_list = []
    
    def __getitem__(self, index):
        '''
        返回一条样本
        '''
        if self.keep_in_memory:
            data = self.data
            prompt, response = data.iloc[index].prompt, data.iloc[index].response
        else:
            prompt, response = next(self.sample_generator)

        max_seq_len = self.max_seq_len - 5 # len('[EOS]') = 5
        # add an eos token note that end of resopnse, using in generate.
        return f"{prompt[0: max_seq_len]}[EOS]", f"{response[0: max_seq_len]}[EOS]"

    def collate_fn(self, data: list[list]) -> dict:
        '''
        合并一个批次数据返回
        '''
        tokenizer = self.tokenizer

        prompt = tokenizer([item[0] for item in data], padding=True, return_token_type_ids=False)
        response = tokenizer([item[1] for item in data], padding=True, return_token_type_ids=False)

        input_ids = array(prompt.input_ids, dtype=int64)
        input_mask = array(prompt.attention_mask, dtype=int64)
        target_ids = array(response.input_ids, dtype=int64)

        ret = {
            'input_ids': LongTensor(input_ids),
            'input_mask': LongTensor(input_mask),
            'target_ids': LongTensor(target_ids),
        }
        return ret
    
    def __len__(self) -> int:
        return self.length

class ParquetDataset:
 
    def __init__(self,  
                parquet_file: Union[str, dict],
                tokenizer_dir: str, 
                keep_in_memory: bool=False,
                cache_dir: str='./.cache',
                buffer_size: int=10240, 
                max_len: int=512, 
                seed: int=23333
            ) -> None:
        '''
        使用huggingface的loaddataset方法加载,
        parquet_file: 单个文件，此时只能使用dataset['train']，
                多个文件请用:parquet_file={'train': 'train.parquet', 'test': 'test.parquet', 'validation': 'validation.parquet'})
                其他用法见：https://huggingface.co/docs/datasets/loading
        keep_in_memory: 是否将parquet文件转换为pandas.DataFrame格式存放到内存
        '''
        self.keep_in_memory = keep_in_memory
        self.len_dict = self.__get_all_parquet_file_size(parquet_file=parquet_file)

        self.max_len = max_len
        self.tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_dir)

        self.tokenizer = self.tokenizer
        
        streaming = False if keep_in_memory else True 
        # streaming=True,否则大数据集OOM
        dataset = load_dataset('parquet', data_files=parquet_file, cache_dir=cache_dir, streaming=streaming) 

        # 这里的batch_size不是训练的batch_size，是传递给precess_batch_func批处理的batch_size
        dataset = dataset.map(self.precess_batch_func, batched=True, batch_size=buffer_size, \
                            remove_columns=['prompt', 'response'], fn_kwargs={'max_len': max_len})

        dataset = dataset.with_format(type="torch")

        if keep_in_memory:
           dataset = dataset.shuffle(seed=seed, keep_in_memory=keep_in_memory)
        else:
            # 只能打乱缓冲区内的数据，不能打乱整个数据集，因此可以将缓存区设置稍微大一些
            dataset = dataset.shuffle(seed=seed, buffer_size=buffer_size)

        self.dataset = dataset
    
    @staticmethod
    def precess_batch_func(item: dict, max_len: int=512) -> dict:
        '''
        添加EOS
        '''
        max_len -= 5 # len('[EOS]') = 5
        for i in range(len(item['prompt'])):
            item['prompt'][i] = f"{item['prompt'][i][0: max_len]}[EOS]"
        for i in range(len(item['response'])):
            item['response'][i] = f"{item['response'][i][0: max_len]}[EOS]"

        return {
            'prompt': item['prompt'],
            'response': item['response'],
        }
    
    def collate_fn(self, data: list[list]) -> dict:
        '''
        合并一个批次数据返回
        '''
        
        tokenizer = self.tokenizer
        prompt = [item['prompt'] for item in data ]
        response = [item['response'] for item in data ]

        # 按批次pad
        prompt_encoded = tokenizer(prompt, padding=True, return_token_type_ids=False)
        response_encoded = tokenizer(response, padding=True, return_token_type_ids=False)

        input_ids = array(prompt_encoded.input_ids, dtype=int64)
        input_mask = array(prompt_encoded.attention_mask, dtype=int64)
        target_ids = array(response_encoded.input_ids, dtype=int64)

        ret = {
            'input_ids': LongTensor(input_ids),
            'input_mask': LongTensor(input_mask),
            'target_ids': LongTensor(target_ids),
        }
        return ret
    def __getitem__(self, index: str) -> datasets.Dataset:
        '''
        魔术方法，实现下标访问，如：dataset['train']、dataset['validation']、dataset['test']
        '''
        return self.dataset[index]
    
    def __get_all_parquet_file_size(self, parquet_file: Union[str, dict]) -> dict:
        '''
        获取所有parquet file的长度
        '''
        len_dict = dict()
        if type(parquet_file) is str:
            train_len = self.__get_size_of_praquet(parquet_file)
            len_dict['train'] = train_len
        
        if type(parquet_file) is dict:
            for split_type, file in parquet_file.items():
                len_dict[split_type] = self.__get_size_of_praquet(file)
        
        return len_dict
    
    def __get_size_of_praquet(self, file_name: str) -> int:
        '''
        获取一个parquet文件的行数
        '''
        parquet_data = pq.read_table(file_name)

        return parquet_data.num_rows 
    
    def __len__(self) -> int:
        '''
        魔术方法，如果只有一个数据集，返回默认数据集大小
        '''
        if len(self.len_dict) == 1:
            return self.len_dict['train']
        else:
            raise Exception("this dataset contains many splited datasets, use `get_dataset_size(split_name)` function to get length, e.g: get_dataset_size('train')")
    
    def get_dataset_size(self, split_name: str) -> int:
        '''
        获取每个切分数据集的长度
        split_name可取：train、validation、test
        '''
        return self.len_dict[split_name]
    
    def get_tokenizer(self, ) -> PreTrainedTokenizerFast:
        return self.tokenizer



if __name__ == '__main__':
    parquet_file = PROJECT_ROOT + '/data/my_valid_dataset.parquet'
    tokenizer_dir = PROJECT_ROOT + '/model_save/tokenizer'

    # example 1：
    dataset = MyDataset(parquet_file, tokenizer_dir, keep_in_memory=False, max_seq_len=128)
    print('\nexample 1, dataset size: ', len(dataset))
    dataloader = DataLoader(dataset, batch_size=32, collate_fn=dataset.collate_fn)

    for epoch in range(2):
        print('epoch: {}'.format(epoch))
        for step, batch in enumerate(dataloader):
            x, x_mask, y = batch['input_ids'], batch['input_mask'], batch['target_ids']
            print('step:{}'.format(step), x.shape, x_mask.shape, y.shape)
            if step == 5:
                break

    
    # exit(0)
    # example 2:
    dataset = ParquetDataset(parquet_file, tokenizer_dir, keep_in_memory=True, max_len=32)
    dataloader = DataLoader(dataset['train'], batch_size=32, collate_fn=dataset.collate_fn)
    print('\nexample 2, dataset size: ', dataset.get_dataset_size('train'))

    for epoch in range(2):
        print('epoch: {}'.format(epoch))
        for step, batch in enumerate(dataloader):
            x, x_mask, y = batch['input_ids'], batch['input_mask'], batch['target_ids']
            print('step:{}'.format(step), x.shape, x_mask.shape, y.shape)
            if step == 5:
                break