Tokenize the entire dataset and reshape it into sequences of length seqlen.
(self, ds, tokenizer, seqlen=2048)
| 79 | |
| 80 | class TestDataset(Dataset): |
| 81 | def __init__(self, ds, tokenizer, seqlen=2048): |
| 82 | """Tokenize the entire dataset and reshape it into sequences of length seqlen.""" |
| 83 | |
| 84 | tokenized_ds = tokenizer("\n\n".join(ds['text']), return_tensors='pt') |
| 85 | nsamples = tokenized_ds.input_ids.numel() // seqlen |
| 86 | |
| 87 | input_ids = tokenized_ds.input_ids[0, : nsamples * seqlen] |
| 88 | input_ids = input_ids.reshape(nsamples, seqlen) |
| 89 | attn_mask = tokenized_ds.attention_mask[0, : nsamples * seqlen] |
| 90 | attn_mask = attn_mask.reshape(nsamples, seqlen) |
| 91 | |
| 92 | self.input_ids = input_ids |
| 93 | self.attn_mask = attn_mask |
| 94 | |
| 95 | def __getitem__(self, idx): |
| 96 | return {"input_ids": self.input_ids[idx], "attention_mask": self.attn_mask[idx]} |
nothing calls this directly
no outgoing calls
no test coverage detected