| 23 | |
| 24 | |
| 25 | class SFTDataset(Dataset): |
| 26 | def __init__(self, data_dir, tokenizer, data_type='train'): |
| 27 | super().__init__() |
| 28 | |
| 29 | self.data_dir = data_dir |
| 30 | self.tokenizer = tokenizer |
| 31 | self.data_type = data_type |
| 32 | |
| 33 | self.data = [] |
| 34 | # We do not calculate losses for the meta instruction or results returned by plugins |
| 35 | # The token spans with label -100, [(span_start, span_end), ...] |
| 36 | self.no_loss_spans = [] |
| 37 | |
| 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) |