A dataset class that can read data with raw mds-format (mosaic streaming-format without compression) from local. In comparison with `StreamingTextDataset` that also can read data with mds-format from local, this class is slimmer, more efficient, and does not contain redundant code requi
| 445 | |
| 446 | |
| 447 | class NoStreamingDataset(Dataset): |
| 448 | """ |
| 449 | A dataset class that can read data with raw mds-format (mosaic streaming-format without compression) |
| 450 | from local. In comparison with `StreamingTextDataset` that also can read data with mds-format from local, |
| 451 | this class is slimmer, more efficient, and does not contain redundant code required for streaming. |
| 452 | """ |
| 453 | |
| 454 | def __init__( |
| 455 | self, |
| 456 | local: str, |
| 457 | split: Optional[str], |
| 458 | max_seq_len: int, |
| 459 | tokenizer: Optional[Tokenizer] = None, |
| 460 | pad_sequences: bool = True, |
| 461 | ) -> None: |
| 462 | super().__init__() |
| 463 | if split is not None: |
| 464 | split_path = os.path.join(local, split) |
| 465 | else: |
| 466 | split_path = local |
| 467 | index_file_path = os.path.join(split_path, "index.json") |
| 468 | obj = json.load(open(index_file_path)) |
| 469 | self.shards = [] |
| 470 | for info in obj["shards"]: |
| 471 | shard = reader_from_json(local, split, info) |
| 472 | raw_filename = os.path.join(shard.dirname, shard.split, shard.raw_data.basename) |
| 473 | assert os.path.isfile(raw_filename), f"Raw file {raw_filename} does not exist" |
| 474 | shard.validate(True) |
| 475 | self.shards.append(shard) |
| 476 | samples_per_shard = np.array([shard.samples for shard in self.shards], np.int64) |
| 477 | self.len = samples_per_shard.sum() |
| 478 | self.spanner = Spanner(samples_per_shard) |
| 479 | self.max_seq_len = max_seq_len |
| 480 | self.tokenizer = tokenizer |
| 481 | self.pad_sequences = pad_sequences |
| 482 | |
| 483 | def _tokenize(self, text_sample): |
| 484 | assert self.tokenizer is not None, "Tokenizer required if data is not pretokenized" |
| 485 | if self.tokenizer._pad_token is None: |
| 486 | # Some tokenizers (e.g. GPT2 tokenizer) have no padding token which causes bugs |
| 487 | raise RuntimeError("If tokenizing on-the-fly, tokenizer must have a pad_token_id") |
| 488 | |
| 489 | return self.tokenizer( |
| 490 | text_sample["text"], |
| 491 | truncation=True, |
| 492 | padding="max_length" if self.pad_sequences else False, |
| 493 | max_length=self.max_seq_len, |
| 494 | ) |
| 495 | |
| 496 | def __getitem__(self, index: int): |
| 497 | shard_id, shard_sample_id = self.spanner[index] |
| 498 | shard = self.shards[shard_id] |
| 499 | sample = shard[shard_sample_id] |
| 500 | if "input_ids" in sample: |
| 501 | for k in list(sample.keys()): |
| 502 | if isinstance(sample[k], np.ndarray): |
| 503 | sample[k] = sample[k][: self.max_seq_len] |
| 504 | else: |
no outgoing calls
no test coverage detected