File size: 870 Bytes
01f199c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | from datasets.Dataset import Dataset
from datasets.MBPPDataset import MBPPDataset
from datasets.APPSDataset import APPSDataset
from datasets.XCodeDataset import XCodeDataset
from datasets.HumanEvalDataset import HumanDataset
from datasets.CodeContestDataset import CodeContestDataset
class DatasetFactory:
@staticmethod
def get_dataset_class(dataset_name):
if dataset_name == "APPS":
return APPSDataset
elif dataset_name == "MBPP":
return MBPPDataset
elif dataset_name == "XCode":
return XCodeDataset
elif dataset_name == "HumanEval":
return HumanDataset
elif dataset_name == "Human":
return HumanDataset
elif dataset_name == "CC":
return CodeContestDataset
else:
raise Exception(f"Unknown dataset name {dataset_name}")
|