Fraser/dream-coder

1"""Program Synthesis dataset from dreamcoder. https://github.com/ellisk42/ec"""2from random import choice, shuffle3import datasets4import pandas as pd5 6from dreamcoder.domains.text.makeTextTasks import makeTasks as textMakeTasks7from dreamcoder.domains.list.main import main as listMakeTasks8 9 10_DESCRIPTION = """\11Generated program synthesis datasets used to train dreamcoder.12"""13_FEATURES = datasets.Features(14    {15        "description": datasets.Value("string"),16        "input": datasets.Value("string"),17        "output": datasets.Value("string"),18        "types": datasets.Value("string")19    }20)21 22_HOMEPAGE = "https://github.com/ellisk42/ec"23 24_LICENSE = "MIT License"25 26_MAX_STEPS = 378227 28 29class infIterator:30    def __init__(self, make_mthd):31        self.make_mthd = make_mthd32        self.i = None33 34    def reset(self):35        tasks = self.make_mthd()36 37        rows = []38        for task in tasks:39            base = {40                'types': str(task.request),41                "description": task.name,42            }43            for (inp, outp) in task.examples:44                rows.append(dict(input=str(inp), output=str(outp), **base))45 46        shuffle(rows)47        self.rows = rows48        self.i = 049 50    def step(self):51        if self.i is None:52            self.reset()53        row = self.rows[self.i]54        self.i += 155        if self.i >= len(self.rows):56            self.reset()57        return row58 59 60class ProgramSynthesis(datasets.GeneratorBasedBuilder):61    """Program Synthesis dataset from dreamcoder."""62 63    VERSION = datasets.Version("1.1.0")64    BUILDER_CONFIGS = [65        datasets.BuilderConfig(name="text", version=VERSION, description="Text tasks."),66        datasets.BuilderConfig(name="list", version=VERSION, description="List tasks."),67        datasets.BuilderConfig(name="all", version=VERSION, description="All tasks at once."),68    ]69    DEFAULT_CONFIG_NAME = "all"70 71    def _info(self):72        return datasets.DatasetInfo(73            description=_DESCRIPTION,74            features=_FEATURES,75            supervised_keys=("input", "output"),76            homepage=_HOMEPAGE,77            license=_LICENSE,78        )79 80    def _split_generators(self, dl_manager):81        return [82            datasets.SplitGenerator(83                name=datasets.Split.TRAIN, gen_kwargs={'split': 'train'}84            ),85            datasets.SplitGenerator(86                name=datasets.Split.TEST, gen_kwargs={'split': 'test'}87            ),88        ]89 90    def _generate_examples(self, split):91        if split == 'test':92            # TODO read from: test.list.csv, test.text.csv93            df_list = pd.read_csv('_t.list.csv')94            df_text = pd.read_csv('_t.text.csv')95            if self.config.name == 'all':96                df = pd.concat(df_list, df_text)97            elif self.config.name == 'list':98                df = df_list99            elif self.config.name == 'text':100                df = df_text101            else:102                raise Exception('Bad Config')103            for i, row in df.iterrows():104                yield i, dict(row)105            return106 107        task_samples = {108            'text': infIterator(textMakeTasks),109            'list': infIterator(listMakeTasks),110        }111        ks = list(task_samples.keys())112        for key in range(_MAX_STEPS):113 114            if self.config.name == 'all':115                dataset_type = choice(ks)116            else:117                dataset_type = self.config.name118 119            yield key, task_samples[dataset_type].step()120