amachine.am_transformers.am_text_datasets
1import pyarrow as pa 2import pyarrow.parquet as pq 3import pyarrow.types as patypes 4import numpy as np 5from tokenizers import Tokenizer 6from .am_datasets import StreamingParquetDataset 7 8class TextStreamingParquetDataset(StreamingParquetDataset): 9 10 def __init__( 11 self, path: str, 12 input_column: str, 13 tokenizer_path: str, 14 seq_len: int, 15 eos_id: int, 16 end_docs_with_eos: bool | None = None, 17 pretokenized: bool = True, 18 **kwargs 19 ): 20 if pretokenized and end_docs_with_eos is not None: 21 raise ValueError( 22 "end_docs_with_eos has no effect when pretokenized=True: " 23 "pretokenized data is already a flattened 1D stream with no " 24 "recoverable document boundaries" 25 ) 26 27 super().__init__(path=path, input_column=input_column, remap=np.array([]), seq_len=seq_len, **kwargs) 28 self.tokenizer_path = tokenizer_path 29 self.eos_id = eos_id 30 self._tokenizer: Tokenizer | None = None 31 self.end_docs_with_eos = end_docs_with_eos 32 self.pretokenized = pretokenized 33 34 @property 35 def tokenizer(self) -> Tokenizer: 36 if self._tokenizer is None: 37 self._tokenizer = Tokenizer.from_file(self.tokenizer_path) 38 return self._tokenizer 39 40 def _read_row_group(self, parquet: pq.ParquetFile, rg_idx: int) -> dict[str, np.ndarray] | None: 41 42 table = parquet.read_row_group(rg_idx, columns=self.columns) 43 44 if table.num_rows == 0: 45 return None 46 47 column = table.column(self.input_col) 48 49 if self.pretokenized : 50 51 # Catch if the Parquet file has arrays/lists of tokens per row instead of flat ints 52 if patypes.is_list(column.type) or patypes.is_large_list(column.type): 53 raise TypeError( 54 f"Column '{self.input_col}' contains nested lists ({column.type}). " 55 f"When `pretokenized=True`, the dataloader expects a fully flattened " 56 f"1D continuous stream of integers." 57 ) 58 59 # Verify the column is actually integers 60 if not patypes.is_integer(column.type): 61 raise TypeError( 62 f"Column '{self.input_col}' has type {column.type}. " 63 f"Expected a flat integer type." 64 ) 65 66 # Data is correctly flattened ints 67 tokens = column.to_numpy() 68 return {self.input_col: tokens.astype(np.int64)} 69 70 # Not pretokenized: expect strings to encode 71 if not (patypes.is_string(column.type) or patypes.is_large_string(column.type)): 72 raise TypeError( 73 f"Dataset configured as `pretokenized=False`, but Parquet column " 74 f"'{self.input_col}' has type {column.type}. Expected strings." 75 ) 76 77 data_list = column.to_pylist() 78 if not data_list: 79 return None 80 81 all_tokens = [] 82 encoded_batches = self.tokenizer.encode_batch(data_list) 83 84 for encoding in encoded_batches: 85 86 ids = encoding.ids 87 88 if self.end_docs_with_eos : 89 if not ids or ids[-1] != self.eos_id: 90 ids = ids + [self.eos_id] 91 92 else : 93 end = len(ids) 94 while end > 0 and ids[ end - 1 ] == self.eos_id : 95 end -= 1 96 ids = ids[:end] 97 98 all_tokens.extend(ids) 99 100 return {self.input_col: np.array(all_tokens, dtype=np.int64)}
class
TextStreamingParquetDataset(torch.utils.data.dataset.Dataset[+_T_co], typing.Iterable[+_T_co]):
9class TextStreamingParquetDataset(StreamingParquetDataset): 10 11 def __init__( 12 self, path: str, 13 input_column: str, 14 tokenizer_path: str, 15 seq_len: int, 16 eos_id: int, 17 end_docs_with_eos: bool | None = None, 18 pretokenized: bool = True, 19 **kwargs 20 ): 21 if pretokenized and end_docs_with_eos is not None: 22 raise ValueError( 23 "end_docs_with_eos has no effect when pretokenized=True: " 24 "pretokenized data is already a flattened 1D stream with no " 25 "recoverable document boundaries" 26 ) 27 28 super().__init__(path=path, input_column=input_column, remap=np.array([]), seq_len=seq_len, **kwargs) 29 self.tokenizer_path = tokenizer_path 30 self.eos_id = eos_id 31 self._tokenizer: Tokenizer | None = None 32 self.end_docs_with_eos = end_docs_with_eos 33 self.pretokenized = pretokenized 34 35 @property 36 def tokenizer(self) -> Tokenizer: 37 if self._tokenizer is None: 38 self._tokenizer = Tokenizer.from_file(self.tokenizer_path) 39 return self._tokenizer 40 41 def _read_row_group(self, parquet: pq.ParquetFile, rg_idx: int) -> dict[str, np.ndarray] | None: 42 43 table = parquet.read_row_group(rg_idx, columns=self.columns) 44 45 if table.num_rows == 0: 46 return None 47 48 column = table.column(self.input_col) 49 50 if self.pretokenized : 51 52 # Catch if the Parquet file has arrays/lists of tokens per row instead of flat ints 53 if patypes.is_list(column.type) or patypes.is_large_list(column.type): 54 raise TypeError( 55 f"Column '{self.input_col}' contains nested lists ({column.type}). " 56 f"When `pretokenized=True`, the dataloader expects a fully flattened " 57 f"1D continuous stream of integers." 58 ) 59 60 # Verify the column is actually integers 61 if not patypes.is_integer(column.type): 62 raise TypeError( 63 f"Column '{self.input_col}' has type {column.type}. " 64 f"Expected a flat integer type." 65 ) 66 67 # Data is correctly flattened ints 68 tokens = column.to_numpy() 69 return {self.input_col: tokens.astype(np.int64)} 70 71 # Not pretokenized: expect strings to encode 72 if not (patypes.is_string(column.type) or patypes.is_large_string(column.type)): 73 raise TypeError( 74 f"Dataset configured as `pretokenized=False`, but Parquet column " 75 f"'{self.input_col}' has type {column.type}. Expected strings." 76 ) 77 78 data_list = column.to_pylist() 79 if not data_list: 80 return None 81 82 all_tokens = [] 83 encoded_batches = self.tokenizer.encode_batch(data_list) 84 85 for encoding in encoded_batches: 86 87 ids = encoding.ids 88 89 if self.end_docs_with_eos : 90 if not ids or ids[-1] != self.eos_id: 91 ids = ids + [self.eos_id] 92 93 else : 94 end = len(ids) 95 while end > 0 and ids[ end - 1 ] == self.eos_id : 96 end -= 1 97 ids = ids[:end] 98 99 all_tokens.extend(ids) 100 101 return {self.input_col: np.array(all_tokens, dtype=np.int64)}
An iterable Dataset.
All datasets that represent an iterable of data samples should subclass it. Such form of datasets is particularly useful when data come from a stream.
All subclasses should overwrite __iter__(), which would return an
iterator of samples in this dataset.
When a subclass is used with ~torch.utils.data.DataLoader, each
item in the dataset will be yielded from the ~torch.utils.data.DataLoader
iterator. When num_workers > 0, each worker process will have a
different copy of the dataset object, so it is often desired to configure
each copy independently to avoid having duplicate data returned from the
workers. ~torch.utils.data.get_worker_info(), when called in a worker
process, returns information about the worker. It can be used in either the
dataset's __iter__() method or the ~torch.utils.data.DataLoader 's
worker_init_fn option to modify each copy's behavior.
Example 1: splitting workload across all workers in __iter__()::
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_DATALOADER)
>>> # xdoctest: +SKIP("Fails on MacOS12")
>>> class MyIterableDataset(torch.utils.data.IterableDataset):
... def __init__(self, start, end):
... super(MyIterableDataset).__init__()
... assert end > start, "this example only works with end >= start"
... self.start = start
... self.end = end
...
... def __iter__(self):
... worker_info = torch.utils.data.get_worker_info()
... if worker_info is None: # single-process data loading, return the full iterator
... iter_start = self.start
... iter_end = self.end
... else: # in a worker process
... # split workload
... per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))
... worker_id = worker_info.id
... iter_start = self.start + worker_id * per_worker
... iter_end = min(iter_start + per_worker, self.end)
... return iter(range(iter_start, iter_end))
...
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
>>> ds = MyIterableDataset(start=3, end=7)
>>> # Single-process loading
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
[tensor([3]), tensor([4]), tensor([5]), tensor([6])]
>>> # xdoctest: +REQUIRES(POSIX)
>>> # Multi-process loading with two worker processes
>>> # Worker 0 fetched [3, 4]. Worker 1 fetched [5, 6].
>>> # xdoctest: +IGNORE_WANT("non deterministic")
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
[tensor([3]), tensor([5]), tensor([4]), tensor([6])]
>>> # With even more workers
>>> # xdoctest: +IGNORE_WANT("non deterministic")
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=12)))
[tensor([3]), tensor([5]), tensor([4]), tensor([6])]
Example 2: splitting workload across all workers using worker_init_fn::
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_DATALOADER)
>>> class MyIterableDataset(torch.utils.data.IterableDataset):
... def __init__(self, start, end):
... super(MyIterableDataset).__init__()
... assert end > start, "this example only works with end >= start"
... self.start = start
... self.end = end
...
... def __iter__(self):
... return iter(range(self.start, self.end))
...
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
>>> ds = MyIterableDataset(start=3, end=7)
>>> # Single-process loading
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
[3, 4, 5, 6]
>>>
>>> # Directly doing multi-process loading yields duplicate data
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
[3, 3, 4, 4, 5, 5, 6, 6]
>>> # Define a `worker_init_fn` that configures each dataset copy differently
>>> def worker_init_fn(worker_id):
... worker_info = torch.utils.data.get_worker_info()
... dataset = worker_info.dataset # the dataset copy in this worker process
... overall_start = dataset.start
... overall_end = dataset.end
... # configure the dataset to only process the split workload
... per_worker = int(math.ceil((overall_end - overall_start) / float(worker_info.num_workers)))
... worker_id = worker_info.id
... dataset.start = overall_start + worker_id * per_worker
... dataset.end = min(dataset.start + per_worker, overall_end)
...
>>> # Mult-process loading with the custom `worker_init_fn`
>>> # Worker 0 fetched [3, 4]. Worker 1 fetched [5, 6].
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2, worker_init_fn=worker_init_fn)))
[3, 5, 4, 6]
>>> # With even more workers
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=12, worker_init_fn=worker_init_fn)))
[3, 4, 5, 6]
TextStreamingParquetDataset( path: str, input_column: str, tokenizer_path: str, seq_len: int, eos_id: int, end_docs_with_eos: bool | None = None, pretokenized: bool = True, **kwargs)
11 def __init__( 12 self, path: str, 13 input_column: str, 14 tokenizer_path: str, 15 seq_len: int, 16 eos_id: int, 17 end_docs_with_eos: bool | None = None, 18 pretokenized: bool = True, 19 **kwargs 20 ): 21 if pretokenized and end_docs_with_eos is not None: 22 raise ValueError( 23 "end_docs_with_eos has no effect when pretokenized=True: " 24 "pretokenized data is already a flattened 1D stream with no " 25 "recoverable document boundaries" 26 ) 27 28 super().__init__(path=path, input_column=input_column, remap=np.array([]), seq_len=seq_len, **kwargs) 29 self.tokenizer_path = tokenizer_path 30 self.eos_id = eos_id 31 self._tokenizer: Tokenizer | None = None 32 self.end_docs_with_eos = end_docs_with_eos 33 self.pretokenized = pretokenized