Dataset for supervised fine-tuning.
| 192 | |
| 193 | |
| 194 | class LazySupervisedDataset(Dataset): |
| 195 | """Dataset for supervised fine-tuning.""" |
| 196 | |
| 197 | def __init__(self, raw_data, tokenizer: transformers.PreTrainedTokenizer, template="tool-llama"): |
| 198 | super(LazySupervisedDataset, self).__init__() |
| 199 | self.tokenizer = tokenizer |
| 200 | |
| 201 | rank0_print("Formatting inputs...Skip in lazy mode") |
| 202 | self.tokenizer = tokenizer |
| 203 | self.raw_data = raw_data |
| 204 | self.cached_data_dict = {} |
| 205 | self.template = template |
| 206 | |
| 207 | def __len__(self): |
| 208 | return len(self.raw_data) |
| 209 | |
| 210 | def __getitem__(self, i) -> Dict[str, torch.Tensor]: |
| 211 | if i in self.cached_data_dict: |
| 212 | return self.cached_data_dict[i] |
| 213 | |
| 214 | ret = preprocess([self.raw_data[i]["conversations"]], self.tokenizer, self.template) |
| 215 | ret = dict( |
| 216 | input_ids=ret["input_ids"][0], |
| 217 | labels=ret["labels"][0], |
| 218 | attention_mask=ret["attention_mask"][0], |
| 219 | ) |
| 220 | self.cached_data_dict[i] = ret |
| 221 | |
| 222 | return ret |
| 223 | |
| 224 | |
| 225 | def make_supervised_data_module( |
nothing calls this directly
no outgoing calls
no test coverage detected