| 38 | self.load_data() |
| 39 | |
| 40 | def load_data(self): |
| 41 | logger.info("Loading data...") |
| 42 | data_file = os.path.join(self.data_dir, f'{self.data_type}_data') |
| 43 | no_loss_spans_file = os.path.join(self.data_dir, f'{self.data_type}_no_loss_spans') |
| 44 | if os.path.exists(data_file) and os.path.exists(no_loss_spans_file): |
| 45 | self.data = torch.load(data_file, map_location='cpu') |
| 46 | self.no_loss_spans = torch.load(no_loss_spans_file, map_location='cpu') |
| 47 | else: |
| 48 | with open(os.path.join(self.data_dir, f'{self.data_type}.jsonl'), 'r') as f: |
| 49 | for line in f: |
| 50 | sample = json.loads(line) |
| 51 | |
| 52 | chat = sample['chat'] |
| 53 | num_turns = int(sample['num_turns']) |
| 54 | |
| 55 | meta_instruction = sample['meta_instruction'] |
| 56 | instruction_ids = self.tokenizer.encode(meta_instruction) |
| 57 | assert isinstance(instruction_ids, list) and len(instruction_ids) > 0 |
| 58 | |
| 59 | input_ids = copy.deepcopy(instruction_ids) |
| 60 | no_loss_spans = [(0, len(instruction_ids))] |
| 61 | |
| 62 | for i in range(num_turns): |
| 63 | cur_turn_ids = [] |
| 64 | cur_no_loss_spans = [] |
| 65 | cur_turn = chat[f'turn_{i+1}'] |
| 66 | for key, value in cur_turn.items(): |
| 67 | |
| 68 | cur_ids = self.tokenizer.encode(value) |
| 69 | |
| 70 | if key == 'Tool Responses': |
| 71 | # The format tokens (<|Results|>:...<eor>\n) should have losses. |
| 72 | cur_no_loss_spans.append((len(input_ids + cur_turn_ids) + 5, len(input_ids + cur_turn_ids + cur_ids) - 2)) |
| 73 | |
| 74 | assert isinstance(cur_ids, list) and len(cur_ids) > 0 |
| 75 | |
| 76 | cur_turn_ids.extend(cur_ids) |
| 77 | |
| 78 | if len(input_ids + cur_turn_ids) > 2048: |
| 79 | break |
| 80 | |
| 81 | input_ids.extend(cur_turn_ids) |
| 82 | no_loss_spans.extend(cur_no_loss_spans) |
| 83 | |
| 84 | if len(input_ids) == len(instruction_ids): |
| 85 | continue |
| 86 | |
| 87 | assert len(input_ids) > 0 and len(input_ids) <= 2048 |
| 88 | |
| 89 | self.data.append(input_ids) |
| 90 | self.no_loss_spans.append(no_loss_spans) |
| 91 | |
| 92 | torch.save(self.data, data_file) |
| 93 | torch.save(self.no_loss_spans, no_loss_spans_file) |
| 94 | |
| 95 | logger.info(f"Load data successfully, total {len(self.data)} training samples") |
| 96 | |
| 97 | def __len__(self): |