Dataset for supervised fine-tuning.
| 293 | |
| 294 | |
| 295 | class LazySupervisedDataset(Dataset): |
| 296 | """Dataset for supervised fine-tuning.""" |
| 297 | |
| 298 | def __init__(self, raw_data, tokenizer: transformers.PreTrainedTokenizer, conv_template = "vicuna-1.1", mask_user = True): |
| 299 | super(LazySupervisedDataset, self).__init__() |
| 300 | self.tokenizer = tokenizer |
| 301 | |
| 302 | rank0_print("Formatting inputs...Skip in lazy mode") |
| 303 | self.conv_template = conv_template |
| 304 | self.mask_user = mask_user |
| 305 | self.tokenizer = tokenizer |
| 306 | self.raw_data = raw_data |
| 307 | self.cached_data_dict = {} |
| 308 | |
| 309 | if mask_user: |
| 310 | rank0_print( |
| 311 | f"WARNING: The loss of user prompt will be masked" |
| 312 | ) |
| 313 | else: |
| 314 | rank0_print( |
| 315 | f"WARNING: The loss of user prompt will **NOT** be masked" |
| 316 | ) |
| 317 | |
| 318 | def __len__(self): |
| 319 | return len(self.raw_data) |
| 320 | |
| 321 | def __getitem__(self, i) -> Dict[str, torch.Tensor]: |
| 322 | if i in self.cached_data_dict: |
| 323 | return self.cached_data_dict[i] |
| 324 | |
| 325 | ret = preprocess([self.raw_data[i]["conversations"]], self.tokenizer, self.conv_template, self.mask_user) |
| 326 | ret = dict( |
| 327 | input_ids=ret["input_ids"][0], |
| 328 | labels=ret["labels"][0], |
| 329 | attention_mask=ret["attention_mask"][0], |
| 330 | ) |
| 331 | self.cached_data_dict[i] = ret |
| 332 | return ret |
| 333 | |
| 334 | |
| 335 | def make_supervised_data_module( |
nothing calls this directly
no outgoing calls
no test coverage detected