Starcoder dataset class
| 21 | from linear_moe.tokenizer import get_tokenizer |
| 22 | |
| 23 | class StarcoderRawDataset(torch.utils.data.Dataset): |
| 24 | """Starcoder dataset class""" |
| 25 | def __init__(self, datapath, max_padding_length): |
| 26 | """ |
| 27 | Args: |
| 28 | datapath (str): The path of the dataset. |
| 29 | max_padding_length (int): The maximum length to pad the input sequences to. |
| 30 | """ |
| 31 | |
| 32 | self.IGNORE_INDEX = -100 |
| 33 | self.tokenizer = get_tokenizer() |
| 34 | self.max_padding_length = max_padding_length |
| 35 | PROMPT_DICT = { |
| 36 | 'prompt_input': |
| 37 | ('<|user|>{instruction}{input}\n\n<|bot|>'), |
| 38 | 'prompt_no_input': |
| 39 | ('<|user|>{instruction}\n\n<|bot|>'), |
| 40 | } |
| 41 | |
| 42 | list_data_dict = self.jload(datapath) |
| 43 | prompt_input, prompt_no_input = PROMPT_DICT[ |
| 44 | 'prompt_input'], PROMPT_DICT['prompt_no_input'] |
| 45 | sources = [ |
| 46 | prompt_input.format_map(example) if example.get('input', '') != '' |
| 47 | else prompt_no_input.format_map(example) |
| 48 | for example in list_data_dict |
| 49 | ] |
| 50 | if 'output' in list_data_dict[0].keys(): |
| 51 | temp = 'output' |
| 52 | elif 'content' in list_data_dict[0].keys(): |
| 53 | temp = 'content' |
| 54 | targets = [ |
| 55 | f"{example[temp]}{self.tokenizer.eos_token}" |
| 56 | for example in list_data_dict |
| 57 | ] |
| 58 | data_dict = self.preprocess(sources, targets, self.tokenizer) |
| 59 | |
| 60 | self.input_ids = data_dict['input_ids'] |
| 61 | self.labels = data_dict['labels'] |
| 62 | self.samples = [] |
| 63 | for inputs, labels in zip(self.input_ids, self.labels): |
| 64 | self.samples.append([inputs, labels]) |
| 65 | |
| 66 | print(' >> total number of samples: {}'.format(len(self.samples))) |
| 67 | |
| 68 | def _make_r_io_base(self, f, mode: str): |
| 69 | if not isinstance(f, io.IOBase): |
| 70 | f = open(f, mode=mode, encoding='utf-8') |
| 71 | return f |
| 72 | |
| 73 | def jload(self, f, mode='r'): |
| 74 | """Load a .json file into a dictionary.""" |
| 75 | f = self._make_r_io_base(f, mode) |
| 76 | jdict = json.load(f) |
| 77 | f.close() |
| 78 | return jdict |
| 79 | |
| 80 | def __len__(self): |
no outgoing calls
no test coverage detected