expr.py

   1 #!/usr/bin/env python
   2
   3 # Any copyright is dedicated to the Public Domain.
   4 # https://creativecommons.org/publicdomain/zero/1.0/
   5
   6 # Written by Francois Fleuret <francois@fleuret.org>
   7
   8 import math, re
   9
  10 import torch, torchvision
  11
  12 from torch import nn
  13 from torch.nn import functional as F
  14
  15
  16 def random_var(nb_variables=None, variables=None):
  17     if variables is None:
  18         return chr(ord("A") + torch.randint(nb_variables, (1,)).item())
  19     else:
  20         l = list(variables)
  21         return l[torch.randint(len(l), (1,)).item()]
  22
  23
  24 def random_expr(variables, operand_max, budget):
  25     if budget <= 5:
  26         op = torch.randint(2, (1,)).item()
  27         if op == 0 and len(variables) > 0:
  28             return random_var(variables=variables)
  29         else:
  30             return str(torch.randint(operand_max + 1, (1,)).item())
  31     else:
  32         op = torch.randint(3, (1,)).item()
  33         if op == 0:
  34             e = random_expr(variables, operand_max, budget - 2)
  35             if ("+" in e or "-" in e or "*" in e) and (e[0] != "(" or e[-1] != ")"):
  36                 return "(" + e + ")"
  37             else:
  38                 return e
  39         else:
  40             b = 2 + torch.randint(budget - 5, (1,)).item()
  41             e1 = random_expr(variables, operand_max, b)
  42             e2 = random_expr(variables, operand_max, budget - b - 1)
  43             if op == 1:
  44                 return e1 + "+" + e2
  45             elif op == 2:
  46                 return e1 + "*" + e2
  47
  48
  49 def generate_program(nb_variables, operand_max, length):
  50     s = ""
  51     variables = set()
  52
  53     while len(s) < length:
  54         v = random_var(nb_variables=nb_variables)
  55         s += v + "=" + random_expr(variables, operand_max, budget=20) + ";"
  56         variables.add(v)
  57
  58     return s, variables
  59
  60
  61 def generate_sequences(nb, nb_variables=5, length=20, operand_max=9, result_max=99):
  62     assert nb_variables <= 26
  63     sequences = []
  64
  65     for n in range(nb):
  66         # We take length itself half of the time, and uniform between
  67         # 1 and length otherwise. The actual length can be slightly
  68         # greater
  69
  70         l = min(length, 1 + torch.randint(length * 2, (1,)).item())
  71         result = None
  72         while result == None or max(result.values()) > result_max:
  73             p, v = generate_program(nb_variables, operand_max, l)
  74             v = ", ".join(['"' + v + '": ' + v for v in v])
  75             ldict = {}
  76             exec(p + "result={" + v + "}", globals(), ldict)
  77             result = ldict["result"]
  78
  79         k = list(result.keys())
  80         k.sort()
  81         sequences.append(p + " " + "".join([v + ":" + str(result[v]) + ";" for v in k]))
  82
  83     return sequences
  84
  85
  86 def extract_results(seq):
  87     f = lambda a: (a[0], -1 if a[1] == "" else int(a[1]))
  88     results = [
  89         dict([f(tuple(x.split(":"))) for x in re.findall("[A-Z]:[0-9]*", s)])
  90         for s in seq
  91     ]
  92     return results
  93
  94
  95 if __name__ == "__main__":
  96     import time
  97
  98     start_time = time.perf_counter()
  99     sequences = generate_sequences(1000, length=40)
 100     end_time = time.perf_counter()
 101     for s in sequences[:10]:
 102         print(s)
 103     print(f"{len(sequences) / (end_time - start_time):.02f} samples per second")
 104
 105     print(extract_results(sequences[:10]))