| 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): |
| 49 | # dataset = load_dataset("json", data_files="/root/model/datasets/code/EvolInstruct-Code-80k.json", split="train") |
| 50 | dataset = load_dataset("json", data_files="nickrosh/Evol-Instruct-Code-80k-v1", split="train") |
| 51 | |
| 52 | dataset = dataset.shuffle(seed=42) |
| 53 | samples = [] |
| 54 | n_run = 0 |
| 55 | |
| 56 | for data in dataset: |
| 57 | istr = data["instruction"] |
| 58 | opt = data["output"] |
| 59 | line = f"Instruction:\n{istr}\nOutput:\n" |
| 60 | line += opt |
| 61 | line = line.strip() |
| 62 | line_encoded = tokenizer.encode(line) |
| 63 | if len(line_encoded) > 512: |
| 64 | continue |
| 65 | sample = torch.tensor([line_encoded]) |
| 66 | if sample.numel() == 0: |
| 67 | continue |
| 68 | samples.append(sample) |
| 69 | n_run += 1 |
| 70 | if n_run == n_samples: |
| 71 | break |
| 72 | # now concatenate all samples and split according to block size |
| 73 | cat_samples = torch.cat(samples, dim=1) |
| 74 | n_split = cat_samples.shape[1] // block_size |
| 75 | print(f" * Split into {n_split} blocks") |
| 76 | return [cat_samples[:, i*block_size:(i+1)*block_size] for i in range(n_split)] |
| 77 | |
| 78 | def get_calib_dataset_gsm8k(tokenizer=None, n_samples=512, block_size=512): |
| 79 | # download from here: https://github.com/OFA-Sys/gsm8k-ScRel/blob/main/data/train_use.jsonl |