Skip to content

Commit

Permalink
read datasets locally
Browse files Browse the repository at this point in the history
  • Loading branch information
TJ-Solergibert committed Sep 17, 2024
1 parent 81fdb3a commit f3bf21d
Showing 1 changed file with 6 additions and 4 deletions.
10 changes: 6 additions & 4 deletions src/nanotron/data/chat_dataset.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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):
"""
Expand Down Expand Up @@ -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)

Expand Down

0 comments on commit f3bf21d

Please sign in to comment.