Fraser/dream-coder

1import datetime2import json3import os4import pickle5import subprocess6import sys7 8from dreamcoder.fragmentGrammar import FragmentGrammar9from dreamcoder.frontier import Frontier, FrontierEntry10from dreamcoder.grammar import Grammar11from dreamcoder.task import Task12from dreamcoder.program import Program, Invented13from dreamcoder.utilities import eprint, timing, callCompiled, get_root_dir14from dreamcoder.vs import induceGrammar_Beta15 16 17def induceGrammar(*args, **kwargs):18    if sum(not f.empty for f in args[1]) == 0:19        eprint("No nonempty frontiers, exiting grammar induction early.")20        return args[0], args[1]21    backend = kwargs.pop("backend", "pypy")22    if 'pypy' in backend:23        # pypy might not like some of the imports needed for the primitives24        # but the primitive values are irrelevant for compression25        # therefore strip them out and then replace them once we are done26        # ditto for task data27        g0,frontiers = args[0].strip_primitive_values(), \28                       [front.strip_primitive_values() for front in args[1]]29        original_tasks = {f.task.name: f.task for f in frontiers}30        frontiers = [Frontier(f.entries, Task(f.task.name,f.task.request,[]))31                     for f in frontiers ]32        args = [g0,frontiers]33 34    35    with timing("Induced a grammar"):36        if backend == "pypy":37            g, newFrontiers = callCompiled(pypyInduce, *args, **kwargs)38        elif backend == "rust":39            g, newFrontiers = rustInduce(*args, **kwargs)40        elif backend == "vs":41            g, newFrontiers = rustInduce(*args, vs=True, **kwargs)42        elif backend == "pypy_vs":43            kwargs.pop('iteration')44            kwargs.pop('topk_use_only_likelihood')45            fn = '/tmp/vs.pickle'46            with open(fn, 'wb') as handle:47                pickle.dump((args, kwargs), handle)48            eprint("For debugging purposes, the version space compression invocation has been saved to", fn)49            g, newFrontiers = callCompiled(induceGrammar_Beta, *args, **kwargs)50        elif backend == "ocaml":51            kwargs.pop('iteration')52            kwargs.pop('topk_use_only_likelihood')53            kwargs['topI'] = 30054            kwargs['bs'] = 100000055            g, newFrontiers = ocamlInduce(*args, **kwargs)56        elif backend == "memorize":57            g, newFrontiers = memorizeInduce(*args, **kwargs)58        else:59            assert False, "unknown compressor"60 61    if 'pypy' in backend:62        g, newFrontiers = g.unstrip_primitive_values(), \63                          [front.unstrip_primitive_values() for front in newFrontiers]64        newFrontiers = [Frontier(f.entries, original_tasks[f.task.name])65                        for f in newFrontiers] 66        67 68    return g, newFrontiers69 70def memorizeInduce(g, frontiers, **kwargs):71    existingInventions = {p.uncurry()72                          for p in g.primitives }73    programs = {f.bestPosterior.program for f in frontiers if not f.empty}74    newInventions = programs - existingInventions75    newGrammar = Grammar.uniform([p for p in g.primitives] + \76                                 [Invented(ni) for ni in newInventions])77    78    # rewrite in terms of new primitives79    def substitute(p):80        nonlocal newInventions81        if p in newInventions: return Invented(p).uncurry()82        return p83    newFrontiers = [Frontier([FrontierEntry(program=np,84                                            logPrior=newGrammar.logLikelihood(f.task.request, np),85                                            logLikelihood=e.logLikelihood)86                           for e in f87                           for np in [substitute(e.program)] ],88                             task=f.task)89                 for f in frontiers ]90    return newGrammar, newFrontiers91    92    93        94    95 96def pypyInduce(*args, **kwargs):97    kwargs.pop('iteration')98    return FragmentGrammar.induceFromFrontiers(*args, **kwargs)99 100 101def ocamlInduce(g, frontiers, _=None,102                topK=1, pseudoCounts=1.0, aic=1.0,103                structurePenalty=0.001, a=0, CPUs=1,104                bs=1000000, topI=300):105    # This is a dirty hack!106    # Memory consumption increases with the number of CPUs107    # And early on we have a lot of stuff to compress108    # If this is the first iteration, only use a fraction of the available CPUs109    if all(not p.isInvented for p in g.primitives):110        if a > 3:111            CPUs = max(1, int(CPUs / 6))112        else:113            CPUs = max(1, int(CPUs / 3))114    else:115        CPUs = max(1, int(CPUs / 2))116    CPUs = 2117 118    # X X X FIXME X X X119    # for unknown reasons doing compression all in one go works correctly and doing it with Python and the outer loop causes problems120    iterations = 99  # maximum number of components to add at once121 122    while True:123        g0 = g124 125        originalFrontiers = frontiers126        t2f = {f.task: f for f in frontiers}127        frontiers = [f for f in frontiers if not f.empty]128        message = {"arity": a,129                   "topK": topK,130                   "pseudoCounts": float(pseudoCounts),131                   "aic": aic,132                   "bs": bs,133                   "topI": topI,134                   "structurePenalty": float(structurePenalty),135                   "CPUs": CPUs,136                   "DSL": g.json(),137                   "iterations": iterations,138                   "frontiers": [f.json()139                                 for f in frontiers]}140 141        message = json.dumps(message)142        if True:143            timestamp = datetime.datetime.now().isoformat()144            os.system("mkdir  -p compressionMessages")145            fn = "compressionMessages/%s" % timestamp146            with open(fn, "w") as f:147                f.write(message)148            eprint("Compression message saved to:", fn)149 150        try:151            # Get relative path152            compressor_file = os.path.join(get_root_dir(), 'compression')153            process = subprocess.Popen(compressor_file,154                                       stdin=subprocess.PIPE,155                                       stdout=subprocess.PIPE)156            response, error = process.communicate(bytes(message, encoding="utf-8"))157            response = json.loads(response.decode("utf-8"))158        except OSError as exc:159            raise exc160 161        g = response["DSL"]162        g = Grammar(g["logVariable"],163                    [(l, p.infer(), p)164                     for production in g["productions"]165                     for l in [production["logProbability"]]166                     for p in [Program.parse(production["expression"])]],167                    continuationType=g0.continuationType)168 169        frontiers = {original.task:170                         Frontier([FrontierEntry(p,171                                                 logLikelihood=e["logLikelihood"],172                                                 logPrior=g.logLikelihood(original.task.request, p))173                                   for e in new["programs"]174                                   for p in [Program.parse(e["program"])]],175                                  task=original.task)176                     for original, new in zip(frontiers, response["frontiers"])}177        frontiers = [frontiers.get(f.task, t2f[f.task])178                     for f in originalFrontiers]179        if iterations == 1 and len(g) > len(g0):180            eprint("Grammar changed - running another round of consolidation.")181            continue182        else:183            eprint("Finished consolidation.")184            return g, frontiers185 186 187def rustInduce(g0, frontiers, _=None,188               topK=1, pseudoCounts=1.0, aic=1.0,189               structurePenalty=0.001, a=0, CPUs=1, iteration=-1,190               topk_use_only_likelihood=False,191               vs=False):192    def finite_logp(l):193        return l if l != float("-inf") else -1000194 195    message = {196        "strategy": {"version-spaces": {"top_i": 50}}197        if vs else198        {"fragment-grammars": {}},199        "params": {200            "structure_penalty": structurePenalty,201            "pseudocounts": int(pseudoCounts + 0.5),202            "topk": topK,203            "topk_use_only_likelihood": topk_use_only_likelihood,204            "aic": aic if aic != float("inf") else None,205            "arity": a,206        },207        "primitives": [{"name": p.name, "tp": str(t), "logp": finite_logp(l)}208                       for l, t, p in g0.productions if p.isPrimitive],209        "inventions": [{"expression": str(p.body),210                        "logp": finite_logp(l)}  # -inf=-100211                       for l, t, p in g0.productions if p.isInvented],212        "variable_logprob": finite_logp(g0.logVariable),213        "frontiers": [{214            "task_tp": str(f.task.request),215            "solutions": [{216                "expression": str(e.program),217                "logprior": finite_logp(e.logPrior),218                "loglikelihood": e.logLikelihood,219            } for e in f],220        } for f in frontiers],221    }222 223    eprint("running rust compressor")224 225    messageJson = json.dumps(message)226 227    with open("jsonDebug", "w") as f:228        f.write(messageJson)229 230    # check which version of python we are using231    # if >=3.6 do:232    if sys.version_info[1] >= 6:233        p = subprocess.Popen(234            ['./rust_compressor/rust_compressor'],235            encoding='utf-8',236            stdin=subprocess.PIPE,237            stdout=subprocess.PIPE)238    elif sys.version_info[1] == 5:239        p = subprocess.Popen(240            ['./rust_compressor/rust_compressor'],241            stdin=subprocess.PIPE,242            stdout=subprocess.PIPE)243 244        messageJson = bytearray(messageJson, encoding='utf-8')245        # convert messageJson string to bytes246    else:247        eprint("must be python 3.5 or 3.6")248        assert False249 250    p.stdin.write(messageJson)251    p.stdin.flush()252    p.stdin.close()253 254    if p.returncode is not None:255        raise ValueError("rust compressor failed")256 257    if sys.version_info[1] >= 6:258        resp = json.load(p.stdout)259    elif sys.version_info[1] == 5:260        import codecs261        resp = json.load(codecs.getreader('utf-8')(p.stdout))262 263    productions = [(x["logp"], p) for p, x in264                   zip((p for (_, _, p) in g0.productions if p.isPrimitive), resp["primitives"])] + \265                  [(i["logp"], Invented(Program.parse(i["expression"])))266                   for i in resp["inventions"]]267    productions = [(l if l is not None else float("-inf"), p)268                   for l, p in productions]269    g = Grammar.fromProductions(productions, resp["variable_logprob"], continuationType=g0.continuationType)270    newFrontiers = [271        Frontier(272            [273                FrontierEntry(274                    Program.parse(275                        s["expression"]),276                    logPrior=s["logprior"],277                    logLikelihood=s["loglikelihood"]) for s in r["solutions"]],278            f.task) for f,279                        r in zip(280            frontiers,281            resp["frontiers"])]282    return g, newFrontiers283