From f3bf21dfff94fd7e68da992b30ce5917887d800f Mon Sep 17 00:00:00 2001 From: tj-solergibert Date: Tue, 17 Sep 2024 10:36:47 +0200 Subject: [PATCH] read datasets locally --- src/nanotron/data/chat_dataset.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/nanotron/data/chat_dataset.py b/src/nanotron/data/chat_dataset.py index be62af6d..df5bd726 100644 --- a/src/nanotron/data/chat_dataset.py +++ b/src/nanotron/data/chat_dataset.py @@ -1,6 +1,9 @@ +import os from typing import List import numpy as np +from datasets import load_dataset +from datasets.distributed import split_dataset_by_node from nanotron.data.chat_tokenizer import ChatTokenizer from nanotron.data.collator import ( build_labels, @@ -11,9 +14,6 @@ from torch.utils.data import IterableDataset from transformers import AutoTokenizer -from datasets import load_dataset -from datasets.distributed import split_dataset_by_node - class ChatDataset(IterableDataset): """ @@ -62,7 +62,9 @@ def __init__( self.seed = seed # Load, split and shuffle dataset - self.dataset = load_dataset(dataset_path, split=split, streaming=True) + self.dataset = load_dataset( + "parquet", split=split, streaming=True, data_files=os.path.join(dataset_path, "data", "*.parquet") + ) self.dataset = split_dataset_by_node(self.dataset, dp_rank, dp_ranks_size) self.dataset = self.dataset.shuffle(seed=seed, buffer_size=10_000)