| 139 | """ |
| 140 | |
| 141 | def __init__( |
| 142 | self, |
| 143 | tokenizer: Tokenizer, |
| 144 | max_seq_len: int, |
| 145 | streams: Optional[Sequence[Stream]] = None, |
| 146 | remote: Optional[str] = None, |
| 147 | local: Optional[str] = None, |
| 148 | split: Optional[str] = None, |
| 149 | download_retry: int = 2, |
| 150 | download_timeout: float = 60, |
| 151 | validate_hash: Optional[str] = None, |
| 152 | keep_zip: bool = False, |
| 153 | epoch_size: Optional[int] = None, |
| 154 | predownload: int = 100_000, |
| 155 | partition_algo: str = "orig", |
| 156 | num_canonical_nodes: Optional[int] = None, |
| 157 | batch_size: Optional[int] = None, |
| 158 | shuffle: bool = False, |
| 159 | shuffle_algo: str = "py1s", |
| 160 | shuffle_seed: int = 9176, |
| 161 | cache_limit: Optional[int] = None, |
| 162 | **kwargs: Dict[str, Any], |
| 163 | ): |
| 164 | group_method = kwargs.pop("group_method", None) |
| 165 | if group_method is not None: |
| 166 | raise NotImplementedError( |
| 167 | "group_method is deprecated and has been removed.\nTo " |
| 168 | + "concatenate, use the --concat_tokens " |
| 169 | + "argument when creating your MDS dataset with concat_c4.py" |
| 170 | ) |
| 171 | |
| 172 | if kwargs is not None and len(kwargs) > 0: |
| 173 | raise ValueError(f"StreamingTextDataset() got an unexpected keyword argument: {kwargs}") |
| 174 | |
| 175 | if local is not None and (remote is None or (local == remote)): |
| 176 | if os.path.isdir(local): |
| 177 | contents = set(os.listdir(local)) |
| 178 | if split not in contents: |
| 179 | raise ValueError(f"local directory {local} does not contain split {split}") |
| 180 | |
| 181 | # Build Dataset |
| 182 | super().__init__( |
| 183 | streams=streams, |
| 184 | remote=remote, |
| 185 | local=local, |
| 186 | split=split, |
| 187 | download_retry=download_retry, |
| 188 | download_timeout=download_timeout, |
| 189 | validate_hash=validate_hash, |
| 190 | keep_zip=keep_zip, |
| 191 | epoch_size=epoch_size, |
| 192 | predownload=predownload, |
| 193 | partition_algo=partition_algo, |
| 194 | num_canonical_nodes=num_canonical_nodes, |
| 195 | batch_size=batch_size, |
| 196 | shuffle=shuffle, |
| 197 | shuffle_algo=shuffle_algo, |
| 198 | shuffle_seed=shuffle_seed, |