Class SupervisedDataset

test/gsm8k/test.py:106–137

Dataset for supervised fine-tuning.

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

    def __init__(self, data_path: str, tokenizer: transformers.PreTrainedTokenizer):
        super(SupervisedDataset, self).__init__()

        # dataset_for_eval = load_dataset(data_path)['train']


        with open(data_path, 'r') as f:
            dataset_for_eval = f.readlines()

        dataset_for_eval = [json.loads(item.strip()) for item in dataset_for_eval]
        try:
            sources = [PROMPT_DICT["prompt_no_input"].format_map(item)for item in dataset_for_eval]
        except:
            sources = [PROMPT_DICT["prompt_no_input_v2"].format_map(item)for item in dataset_for_eval]
        try:
            targets = [item['answer'] for item in dataset_for_eval]
        except:
            targets = [item['response'] for item in dataset_for_eval]

        data_dict = preprocess(sources, targets, tokenizer)

        self.input_ids = data_dict["input_ids"]
        self.labels = data_dict["labels"]

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

    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        return dict(input_ids=self.input_ids[i], labels=self.labels[i], id=i)

def padding(inputs, padding_token, cutoff = None):
    num_elems = len(inputs)