Class SupervisedDataset

train/train.py:175–212

Dataset for supervised fine-tuning.

class SupervisedDataset(Dataset):
    """Dataset for supervised fine-tuning."""

    def __init__(self, data_path: str, tokenizer: transformers.PreTrainedTokenizer, max_sample: int, split: str):
        super().__init__()

        with open(data_path, 'r') as f:
            lines = f.readlines()
        all_dataset = [json.loads(line.strip()) for line in lines]

        sources, targets = zip(*[(s[0][0], f"{s[0][1]}{tokenizer.eos_token}") for s in all_dataset])

        dataset_size = len(sources)
        max_sample = min(max_sample or dataset_size, dataset_size)
        if max_sample < dataset_size:
            indices = random.sample(range(dataset_size), max_sample)
            self.sources, self.targets = [sources[i] for i in indices], [targets[i] for i in indices]
        else:
            self.sources, self.targets = sources, targets 
                 
        split_num = len(self.sources) // 5
        if split == "train":
            self.sources, self.targets = self.sources[split_num:], self.targets[split_num:]
            print(f"Using {len(self.sources)} samples to train")

            print("Example Data")
            print("sources: \n", self.sources[0])
            print("targets: \n", self.targets[0])

        elif split == "eval":
            self.sources, self.targets = self.sources[:split_num], self.targets[:split_num]
            print(f"Using {len(self.sources)} samples to evaluation")

    def __len__(self):
        return len(self.sources)

    def __getitem__(self, i):
        return dict(input_ids=self.sources[i], labels=self.targets[i])

@dataclass
class DataCollatorForSupervisedDataset(object):