| 16 | return get_calib_dataset_code(tokenizer=tokenizer, n_samples=n_samples, block_size=block_size) |
| 17 | |
| 18 | def get_pile_dataset(tokenizer=None, n_samples=512, block_size=512): |
| 19 | # dataset = load_dataset("json", data_files="/root/model/llm-awq/val.jsonl.zst", split="train") |
| 20 | dataset = load_dataset("mit-han-lab/pile-val-backup", split="validation") |
| 21 | dataset = dataset.map(lambda x: { |
| 22 | 'text': x['text'] |
| 23 | }) |
| 24 | dataset = dataset.shuffle(seed=42) |
| 25 | samples = [] |
| 26 | n_run = 0 |
| 27 | |
| 28 | for data in dataset: |
| 29 | line = data["text"] |
| 30 | line = line.strip() |
| 31 | line_encoded = tokenizer.encode(line) |
| 32 | if len(line_encoded) > 512: |
| 33 | continue |
| 34 | sample = torch.tensor([line_encoded]) |
| 35 | if sample.numel() == 0: |
| 36 | continue |
| 37 | samples.append(sample) |
| 38 | n_run += 1 |
| 39 | if n_run == n_samples: |
| 40 | break |
| 41 | # now concatenate all samples and split according to block size |
| 42 | cat_samples = torch.cat(samples, dim=1) |
| 43 | n_split = cat_samples.shape[1] // block_size |
| 44 | print(f" * Split into {n_split} blocks") |
| 45 | return [cat_samples[:, i*block_size:(i+1)*block_size] for i in range(n_split)] |
| 46 | |
| 47 | # TODO: Don't do spliting when code and gsm8k |
| 48 | def get_calib_dataset_code(tokenizer=None, n_samples=512, block_size=512): |