Function make_supervised_data_module

data/generation/single_generate.py:143–147

Make dataset and collator for supervised fine-tuning.

(tokenizer: transformers.PreTrainedTokenizer, dataset_name, max_sample=None)

        )

def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer, dataset_name, max_sample=None) -> Dict:
    """Make dataset and collator for supervised fine-tuning."""
    gen_dataset = SupervisedDataset(tokenizer=tokenizer, dataset_name=dataset_name, max_sample=max_sample)
    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
    return gen_dataset, data_collator

def main(args):
    torch.manual_seed(args.seed)