├── .github
├── ISSUE_TEMPLATE
│ └── new-puzzle.md
└── workflows
│ └── codeql-analysis.yml
├── .gitignore
├── CODE_OF_CONDUCT.md
├── DATASHEET.md
├── ICLR2023
├── README.md
├── data
│ ├── 125M_PAPER_1M_iter_1.txt.gz
│ ├── 13B_PAPER_1M_iter_1.txt.gz
│ ├── 27B_PAPER_1M_iter_1.txt.gz
│ ├── 350M_PAPER_1M_iter_0.txt.gz
│ └── Codex_PAPER_1M_iter_0.txt.gz
├── requirements.txt
└── src
│ ├── babysit.sh
│ ├── ds_config_gptneo.json
│ ├── fine_tune.py
│ ├── fine_tune.sh
│ ├── fine_tune1.sh
│ ├── gen.py
│ ├── gen.sh
│ ├── judge.py
│ ├── neo_train.py
│ ├── preprocess.py
│ ├── requirements.txt
│ ├── solve.py
│ └── utils.py
├── LICENSE
├── README.md
├── SECURITY.md
├── SUPPORT.md
├── generators
├── ICPC.py
├── IMO.py
├── __init__.py
├── algebra.py
├── basic.py
├── chess.py
├── classic_puzzles.py
├── codeforces.py
├── compression.py
├── conways_game_of_life.py
├── games.py
├── graphs.py
├── human_eval.py
├── lattices.py
├── number_theory.py
├── probability.py
├── problem_constants
│ └── rsa_challenges.json
├── study.py
├── trivial_inverse.py
└── tutorial.py
├── make_dataset.py
├── notebooks
├── Demo.ipynb
├── Hackathon_puzzles.ipynb
├── Intro.ipynb
├── demo
│ ├── demo.py
│ └── study_puzzles.py
└── fireworks.gif
├── puzzle_generator.py
├── puzzles
├── README.md
├── puzzles.json
└── split.json
├── solvers
├── README.md
├── codex
│ ├── 138puzzles.json
│ ├── 30puzzles.json
│ ├── 397puzzles.json
│ ├── README.md
│ ├── ezlog.py
│ ├── lm_solve
│ │ ├── __init__.py
│ │ ├── gpt_lib.py
│ │ ├── judge.py
│ │ ├── run.py
│ │ └── scratch.py
│ ├── requirements.txt
│ ├── results
│ │ └── results_397_cushman_codex_1k_full.json.gz
│ ├── run_codex_experiments.py
│ └── utils.py
├── enumerative
│ ├── README.md
│ ├── challenges
│ │ ├── __init__.py
│ │ ├── challenge.py
│ │ └── solutions.py
│ ├── download_pretrained_roberta.sh
│ ├── filter_outputs.py
│ ├── models
│ │ ├── __init__.py
│ │ ├── ml_bow_bigram.py
│ │ ├── model.py
│ │ ├── transformers
│ │ │ ├── __init__.py
│ │ │ ├── dataset_processor.py
│ │ │ ├── finetune_transformer.py
│ │ │ ├── generate_rule_embeddings.py
│ │ │ ├── learn_tokenizer.py
│ │ │ ├── neural_classifier.py
│ │ │ └── preprocess_pretraining_data.py
│ │ └── uniform.py
│ ├── requirements.txt
│ ├── run_bigram.sh
│ ├── run_transformer.sh
│ ├── run_uniform.sh
│ ├── solve_challenges.py
│ ├── top_down.py
│ ├── tython
│ │ ├── __init__.py
│ │ ├── nonterminals.py
│ │ ├── parse.py
│ │ ├── program.py
│ │ └── rules.py
│ └── utils
│ │ ├── __init__.py
│ │ ├── str_utils.py
│ │ └── time_utils.py
└── gpt3
│ ├── README.md
│ ├── create_simple_prompts.py
│ ├── ezlog.py
│ ├── lm_solve
│ ├── __init__.py
│ ├── gpt3_lib.py
│ ├── judge.py
│ └── run.py
│ ├── puzzles_with_descriptions.json
│ ├── puzzles_with_prompts.json
│ ├── requirements.txt
│ ├── run_gpt3_experiments.py
│ └── utils
│ ├── __init__.py
│ ├── str_utils.py
│ └── time_utils.py
└── utils.py
/.github/ISSUE_TEMPLATE/new-puzzle.md:
--------------------------------------------------------------------------------
1 | ---
2 | name: New puzzle
3 | about: Create your own puzzle by defining a function
4 | title: New puzzle
5 | labels: New-puzzle
6 | assignees: ''
7 |
8 | ---
9 |
10 | ```python
11 | def sat(x: str):
12 | """optional problem description"""
13 | return "Hello " + x == "Hello world" # change this to your puzzle
14 | ```
15 |
16 | Solvers, post your solutions in the comments using the following formatting:
17 | ````
18 | Reveal solution
19 |
20 | ```python
21 | def sol():
22 | return "world" # replace with your solution
23 | ```
24 | ", "", "", "", ""],
43 | )
44 |
45 | # Save the files
46 | os.makedirs(args.out, exist_ok=True)
47 | tokenizer.save_model(args.out, args.name)
48 |
49 | # Restoring model from learned vocab/merges
50 | tokenizer = ByteLevelBPETokenizer(
51 | join(args.out, "{}-vocab.json".format(args.name)),
52 | join(args.out, "{}-merges.txt".format(args.name)),
53 | add_prefix_space=True,
54 | )
55 |
56 | # Test encoding
57 | print(tokenizer.encode("Training ByteLevel BPE is very easy").tokens)
58 |
--------------------------------------------------------------------------------
/solvers/enumerative/models/transformers/preprocess_pretraining_data.py:
--------------------------------------------------------------------------------
1 | import glob
2 | import gzip
3 | import json
4 | import random
5 | from tqdm import tqdm
6 |
7 | for split in ['train', 'valid', 'test']:
8 | data_path = 'python_data/{}/*.jsonl.gz'.format(split)
9 | #output_file = 'python_data/{}.jsonl'.format(split)
10 | output_file = 'python_data/{}.txt'.format(split)
11 |
12 | if False:
13 | of = open(output_file, 'w')
14 |
15 | for file_name in glob.glob(data_path):
16 | with gzip.open(file_name, 'rb') as f:
17 | for line in tqdm(f, desc=file_name):
18 | example = json.loads(line)
19 | code = example['code']
20 | doc = example['docstring']
21 | assert doc in code
22 | new_c = code[:code.index(doc)] + code[code.index(doc) + len(doc):]
23 | #new_c = new_c.replace('\n', '\\n')
24 | #doc = doc.replace('\n', '\\n')
25 | if random.random() < 0.5:
26 | inp = new_c
27 | out = doc
28 | else:
29 | inp = doc
30 | out = new_c
31 |
32 | of.write(json.dumps(dict(inp=inp, out=out)) + '\n')
33 |
34 | of.close()
35 |
36 | if True:
37 | of = open(output_file, 'w')
38 |
39 | for file_name in glob.glob(data_path):
40 | with gzip.open(file_name, 'rb') as f:
41 | for line in tqdm(f, desc=file_name):
42 | example = json.loads(line)
43 | of.write(example['original_string'].replace('\n', '\\n'))
44 | of.write('\n')
45 |
46 | of.close()
47 |
48 | if False:
49 | out_splt = split if split != 'valid' else 'val'
50 | src = 'python_data/{}.source'.format(out_splt)
51 | tgt = 'python_data/{}.target'.format(out_splt)
52 | of_src = open(src, 'w')
53 | of_tgt = open(tgt, 'w')
54 | for file_name in glob.glob(data_path):
55 | with gzip.open(file_name, 'rb') as f:
56 | for line in tqdm(f, desc=file_name):
57 | example = json.loads(line)
58 | code = example['code']
59 | doc = example['docstring']
60 | assert doc in code
61 | new_c = code[:code.index(doc)] + code[code.index(doc) + len(doc):]
62 | new_c = new_c.replace('\n', '\\n')
63 | doc = doc.replace('\n', '\\n')
64 | if random.random() < 0.5:
65 | of_src.write(new_c)
66 | of_src.write('\n')
67 | of_tgt.write(doc)
68 | of_tgt.write('\n')
69 | else:
70 | of_src.write(doc)
71 | of_src.write('\n')
72 | of_tgt.write(new_c)
73 | of_tgt.write('\n')
74 | of_src.close()
75 | of_tgt.close()
76 |
--------------------------------------------------------------------------------
/solvers/enumerative/models/uniform.py:
--------------------------------------------------------------------------------
1 | from typing import List, Dict, Callable, Tuple, Generator, Set, Sequence
2 | import functools
3 | import operator
4 | from collections import defaultdict
5 | import random
6 |
7 | from tython import Program, TastNode, _RULES_BY_KIND, Rule, nt
8 | from models.model import CandidateGenerator, reachable_rules_by_kind
9 | from models import RegisterModel
10 | from challenges import extract_constants
11 |
12 |
13 | def prod(iterable): # like sum but product
14 | return functools.reduce(operator.mul, iterable, 1)
15 |
16 |
17 | @RegisterModel("uniform")
18 | class UniformModel(CandidateGenerator):
19 | '''
20 | Uniformly samples from all rules.
21 | '''
22 |
23 | def __init__(self, copy_prob=0.5) -> None:
24 | super().__init__()
25 | self.copy_prob = copy_prob
26 | # self.intended_depth = intended_depth
27 | #self._random_cache = defaultdict(lambda: defaultdict(dict))
28 | self._random_cache = {} # defaultdict is not pickable (for multiproc)
29 |
30 | def random(self, kind, max_depth=5, nodes_by_kind=None):
31 | if nodes_by_kind is None:
32 | nodes_by_kind = {}
33 | key = sum(hash(n) for n in nodes_by_kind)
34 | if key not in self._random_cache:
35 | self._random_cache[key] = {}
36 | cache = self._random_cache[key]
37 |
38 | # cache[kind][depth] is a list of available rules
39 |
40 | def available_rules(kind, depth):
41 | if depth <= 0:
42 | return []
43 | if kind in cache and depth in cache[kind]:
44 | return cache[kind][depth]
45 | rules = [r for r in _RULES_BY_KIND[kind] if
46 | all([nodes_by_kind.get(k) or available_rules(k, depth - 1) for k in r.kids])]
47 | if kind not in cache:
48 | cache[kind] = {}
49 | cache[kind][depth] = rules
50 | return rules
51 |
52 | def helper(kind, depth):
53 | assert depth >= 0
54 | rules = available_rules(kind, depth)
55 | assert rules or nodes_by_kind.get(kind), f"Cannot generate random {kind} of depth <= {depth}"
56 |
57 | if nodes_by_kind.get(kind) and (not rules or random.random() < self.copy_prob):
58 | return random.choice(nodes_by_kind[kind])
59 |
60 | rule = random.choice(rules)
61 |
62 | return TastNode(rule, [helper(k, depth - 1) for k in rule.kids])
63 |
64 | return Program(helper(kind, max_depth))
65 |
66 | def get_candidates_by_nodes(self, kind, nodes_by_kind):
67 | rules_by_kind = _RULES_BY_KIND # zzz reachable_rules_by_kind(kind, nodes_by_kind)
68 |
69 | if kind not in rules_by_kind:
70 | return {}
71 |
72 | # p_kind_rules = {k: (1 - self.copy_prob if nodes_by_kind.get(k) else 1) / max(1, len(rules_by_kind[k]))
73 | # for k in rules_by_kind}
74 |
75 | by_kind = {}
76 |
77 | for parent_kind in rules_by_kind:
78 | rules = rules_by_kind[parent_kind]
79 | has_copy = sum(r.name == "COPY" for r in rules)
80 | if has_copy:
81 | assert has_copy == 1
82 | by_kind[parent_kind] = [], [
83 | (self.copy_prob if r.name == "COPY" else (1 - self.copy_prob) / len(rules), r)
84 | for r in rules]
85 | else:
86 | by_kind[parent_kind] = [], [(1 / len(rules), r) for r in rules]
87 |
88 | ans = {r: [by_kind[k] for k in r.kids]
89 | for rules in rules_by_kind.values() for r in rules if r.name != "COPY"}
90 | ans.update({r: [([(1. / len(nodes_by_kind[r.nt]), n) for n in nodes_by_kind.get(r.nt, [])], [])]
91 | for rules in rules_by_kind.values() for r in rules if r.name == "COPY"})
92 | ans[None] = [by_kind[kind]]
93 |
94 | return ans
95 |
96 | def get_candidates(self, q: Program) -> Dict[Rule,
97 | List[Tuple[List[Tuple[float, TastNode]], List[Tuple[float, Rule]]]]]:
98 |
99 | consts = extract_constants(q)
100 |
101 | sol_type_annotation = q.tree.children[0].children[1].children[1]
102 | sol_kind = nt.type2nt(eval(sol_type_annotation.src()))
103 |
104 | return self.get_candidates_by_nodes(sol_kind, consts)
105 |
106 |
107 | if __name__ == "__main__":
108 | u = UniformModel()
109 | random.seed(0)
110 | for _ in range(100):
111 | p = u.random((List, int), 5)
112 | print(p.src(safe=False))
113 | try:
114 | p.val(max_ticks=1000)
115 | except Program.EvalException:
116 | pass
117 |
--------------------------------------------------------------------------------
/solvers/enumerative/requirements.txt:
--------------------------------------------------------------------------------
1 | tqdm
2 | orderedset
3 | numpy
4 | astor
5 | sklearn
6 | torch==1.13.1
7 | transformers==4.30.0
8 | tensorboardX
9 |
--------------------------------------------------------------------------------
/solvers/enumerative/run_bigram.sh:
--------------------------------------------------------------------------------
1 | #! /bin/bash
2 |
3 | set -ex
4 |
5 | # Bootsrapping process. Starting with Uniform
6 | python solve_challenges.py \
7 | --challenges_path "results/uniform/out.json" \
8 | --eval_challenges_path "results/uniform/out.json" \
9 | -m ml_bow_bigram \
10 | --ml_model rf \
11 | --copy_prob 0.5 \
12 | --max_n_progs 10000 \
13 | --timeout_secs 3600 \
14 | --threads 40 \
15 | --learn_from_train_sols \
16 | --logging_dir results/bootstrap/ml_bigram_0 \
17 | --out_file results/bootstrap/ml_bigram_0/out.json
18 |
19 |
20 | for i in {1..5}; do
21 | python solve_challenges.py \
22 | --challenges_path "results/bootstrap/ml_bigram_$(($i-1))/out.json" \
23 | --eval_challenges_path "results/bootstrap/ml_bigram_$(($i-1))/out.json" \
24 | -m ml_bow_bigram \
25 | --ml_model rf \
26 | --copy_prob 0.5 \
27 | --max_n_progs 10000 \
28 | --timeout_secs 3600 \
29 | --threads 40 \
30 | --learn_from_train_sols \
31 | --logging_dir results/bootstrap/ml_bigram_$i \
32 | --out_file results/bootstrap/ml_bigram_$i/out.json
33 | done
34 |
35 | # Last run until 1M.
36 | python solve_challenges.py \
37 | --challenges_path "results/bootstrap/ml_bigram_5/out.json" \
38 | --eval_challenges_path "results/bootstrap/ml_bigram_5/out.json" \
39 | -m ml_bow_bigram \
40 | --ml_model rf \
41 | --copy_prob 0.5 \
42 | --max_n_progs 1000000 \
43 | --timeout_secs 3600 \
44 | --threads 20 \
45 | --learn_from_train_sols \
46 | --logging_dir results/bootstrap/ml_bigram_6 \
47 | --out_file results/bootstrap/ml_bigram_6/out.json
48 |
49 |
50 | # Run without self-bootrapping (only over unifrom) until 1M.
51 | python solve_challenges.py \
52 | --challenges_path "results/uniform/out.json" \
53 | --eval_challenges_path "results/uniform/out.json" \
54 | -m ml_bow_bigram \
55 | --ml_model rf \
56 | --copy_prob 0.5 \
57 | --max_n_progs 1000000 \
58 | --timeout_secs 3600 \
59 | --threads 40 \
60 | --learn_from_train_sols \
61 | --logging_dir results/ml_bigram \
62 | --out_file results/ml_bigram/out.json
63 |
64 |
--------------------------------------------------------------------------------
/solvers/enumerative/run_transformer.sh:
--------------------------------------------------------------------------------
1 | #! /bin/bash
2 |
3 | set -ex
4 |
5 | # First, extract rule embeddings.
6 | PYTHONPATH=./ python models/transformers/generate_rule_embeddings.py \
7 | --model_name_or_path tals/roberta_python \
8 | --output_dir results/roberta_rule_embeddings
9 |
10 | # Bootsrapping process. Starting with Uniform
11 | PYTHONPATH=./ python models/transformers/finetune_transformer.py \
12 | --challenges_path results/uniform/out.json \
13 | --eval_challenges_path results/uniform/out.json \
14 | --model_name_or_path tals/roberta_python \
15 | --output_dir results/bootstrap/roberta_0 \
16 | --num_train_epochs 20 \
17 | --do_train \
18 | --do_infer \
19 | --max_n_progs 10000 \
20 | --timeout_secs 3600 \
21 | --threads 40 \
22 | --per_gpu_eval_batch_size 128 \
23 | --per_gpu_train_batch_size 16 \
24 | --rule_emb_dir results/roberta_rule_embeddings \
25 | --overwrite_cache \
26 | --max_ticks 10000
27 |
28 | for i in {1..5}; do
29 | PYTHONPATH=./ python models/transformers/finetune_transformer.py \
30 | --challenges_path "results/bootstrap/roberta_$(($i-1))/solutions.json" \
31 | --eval_challenges_path "results/bootstrap/roberta_$(($i-1))/solutions.json" \
32 | --model_name_or_path tals/roberta_python \
33 | --output_dir results/bootstrap/roberta_$i \
34 | --num_train_epochs 20 \
35 | --do_train \
36 | --do_infer \
37 | --max_n_progs 10000 \
38 | --timeout_secs 3600 \
39 | --threads 40 \
40 | --per_gpu_eval_batch_size 128 \
41 | --per_gpu_train_batch_size 16 \
42 | --rule_emb_dir results/roberta_rule_embeddings \
43 | --overwrite_cache \
44 | --max_ticks 10000
45 | done
46 |
47 | # Last run until 1M.
48 | PYTHONPATH=./ python models/transformers/finetune_transformer.py \
49 | --challenges_path "results/bootstrap/roberta_5/solutions.json" \
50 | --eval_challenges_path "results/bootstrap/roberta_5/solutions.json" \
51 | --model_name_or_path tals/roberta_python \
52 | --output_dir results/bootstrap/roberta_6 \
53 | --num_train_epochs 20 \
54 | --do_train \
55 | --do_infer \
56 | --max_n_progs 1000000 \
57 | --timeout_secs 3600 \
58 | --threads 40 \
59 | --per_gpu_eval_batch_size 128 \
60 | --per_gpu_train_batch_size 16 \
61 | --rule_emb_dir results/roberta_rule_embeddings \
62 | --overwrite_cache \
63 | --max_ticks 10000
64 |
65 |
66 | # Run without self-bootrapping (only over unifrom) until 1M.
67 | PYTHONPATH=./ python models/transformers/finetune_transformer.py \
68 | --challenges_path results/uniform/out.json \
69 | --eval_challenges_path results/uniform/out.json \
70 | --model_name_or_path tals/roberta_python \
71 | --output_dir results/roberta \
72 | --num_train_epochs 20 \
73 | --do_train \
74 | --do_infer \
75 | --max_n_progs 1000000 \
76 | --timeout_secs 3600 \
77 | --threads 40 \
78 | --per_gpu_eval_batch_size 128 \
79 | --per_gpu_train_batch_size 16 \
80 | --rule_emb_dir results/roberta_rule_embeddings \
81 | --overwrite_cache \
82 | --max_ticks 10000
83 |
--------------------------------------------------------------------------------
/solvers/enumerative/run_uniform.sh:
--------------------------------------------------------------------------------
1 | python solve_challenges.py \
2 | -p "../../problems/*.json" \
3 | -m uniform \
4 | --solve_uniform \
5 | --copy_prob 0.5 \
6 | --max_n_progs 10000 \
7 | --timeout_secs 3600 \
8 | --threads 40 \
9 | --logging_dir results/uniform \
10 | --out_file results/uniform/out.json
11 |
--------------------------------------------------------------------------------
/solvers/enumerative/tython/__init__.py:
--------------------------------------------------------------------------------
1 | from .program import *
2 | from . import nonterminals as nt
3 | from .rules import Rule, _rules_by_nt as _RULES_BY_KIND, rules as RULES
4 | from .parse import str2name
5 |
--------------------------------------------------------------------------------
/solvers/enumerative/utils/__init__.py:
--------------------------------------------------------------------------------
1 | from utils.str_utils import *
2 | from utils.time_utils import *
3 |
4 | import functools
5 | import operator
6 |
7 |
8 | def prod(iterable): # like sum but product
9 | return functools.reduce(operator.mul, iterable, 1)
10 |
11 |
12 | def flatten(it):
13 | return (e for a in it for e in (flatten(a) if isinstance(a, (tuple, list)) else (a,)))
14 |
15 |
16 | def load_json(filename):
17 | import json
18 | with open(filename, "r") as f:
19 | return json.load(f)
20 |
21 |
22 | def viz_py(py):
23 | import astor, ast
24 | print(astor.dump_tree(ast.parse(py)))
25 |
26 |
27 | def dedup(li):
28 | seen = set()
29 | return [x for x in li if x not in seen and not seen.add(x)]
30 |
--------------------------------------------------------------------------------
/solvers/enumerative/utils/str_utils.py:
--------------------------------------------------------------------------------
1 | import json
2 |
3 | def get_lambda_arg_name(lam):
4 | assert lam.startswith("lambda ")
5 | return lam[len("lambda "):lam.index(":")].strip()
6 |
7 |
8 | def stringify(const):
9 | if type(const) is str:
10 | return json.dumps(const)
11 | return str(const)
12 |
13 |
14 | def color_str(obj, code="\033[0;36m"):
15 | return code + str(obj) + '\033[0m'
16 |
17 |
--------------------------------------------------------------------------------
/solvers/enumerative/utils/time_utils.py:
--------------------------------------------------------------------------------
1 | import time
2 | import logging
3 |
4 | logger = logging.getLogger(__name__)
5 |
6 | def timeit(method):
7 | def timed(*args, **kw):
8 | tick = time.time()
9 | result = method(*args, **kw)
10 | tock = time.time()
11 | logger.debug(f'{method.__name__}: {tock - tick:.3f}s')
12 |
13 | return result
14 | return timed
15 |
--------------------------------------------------------------------------------
/solvers/gpt3/README.md:
--------------------------------------------------------------------------------
1 | # Running GPT-3 experiments
2 |
3 | These are instructions for re-running the GPT-3 experiments. The results will be slightly different than those in
4 | the paper because the API is non-deterministic.
5 |
6 | The requirements can be installed with `pip3 install -r requirements.txt`.
7 |
8 | `run_gpt3_experiments.py` runs the GPT-3 experiments and prints the results to stdout.
9 |
10 | ## Installation and execution.
11 | You will need an open-ai GPT-3 access key which can be signed up for [here](https://openai.com/join/).
12 | You will then need to set it as the `OPENAI_API_KEY` environmnet variable.
13 |
14 | The requirements can be installed with `pip3 install -r requirements.txt`.
15 |
16 | It was run with Python 3.6.9, sys.version = '3.6.9 (default, Jan 26 2021, 15:33:00) \n[GCC 8.4.0]', but should
17 | be compatible with later versions as well.
18 |
19 | Then you simply run
20 | `python run_gpt3_experiments.py` and the results are written to stdout. It uses cacheing mechanisms with the first run
21 | being quite slow and verbose, querying the API. However you can subsequently run it again and it will be
22 | much faster and just output the results. The cacheing makes it deterministic so it should give the same
23 | exact results when re-run.
24 |
--------------------------------------------------------------------------------
/solvers/gpt3/create_simple_prompts.py:
--------------------------------------------------------------------------------
1 | import json
2 | import re
3 |
4 | from run_gpt3_experiments import PARAMS, PREFIX, PREFIX_DOCSTR
5 | import lm_solve
6 |
7 | OUTFILENAME = "puzzles_with_prompts.json"
8 |
9 |
10 | def run(outfilename=OUTFILENAME, filename=PARAMS["filename"]):
11 | with open(filename, "r") as f:
12 | entries = json.load(f)
13 |
14 | for entry in entries:
15 | entry["prompts"] = {}
16 |
17 | for mode in ["short", "medium", "long"]:
18 | prefix = {
19 | "short": "",
20 | "medium": PREFIX,
21 | "long": PREFIX_DOCSTR
22 | }[mode]
23 | prefix = re.sub(r" +$", "", (prefix or "").lstrip(),
24 | flags=re.M) # delete leading/trailing whitespace on each line
25 | puzzles = lm_solve.load_puzzles(filename, mode == "long")
26 | prompts = lm_solve.get_prompts(prefix, [f for ft, f in puzzles])
27 | assert len(puzzles) == len(prompts) == len(entries)
28 | for entry, prompt in zip(entries, prompts):
29 | entry["prompts"][mode] = prompt
30 |
31 | with open(outfilename, "w") as f:
32 | json.dump(entries, f, indent=4)
33 |
34 |
35 | if __name__ == "__main__":
36 | run()
37 |
--------------------------------------------------------------------------------
/solvers/gpt3/ezlog.py:
--------------------------------------------------------------------------------
1 | import os
2 | import logging
3 | import inspect
4 | import io
5 |
6 | my_path = os.path.dirname(__file__)
7 |
8 |
9 | def color_str(obj, code="\033[0;36m"):
10 | return code + str(obj) + '\033[0m'
11 |
12 |
13 | _configured = False
14 |
15 |
16 | def configure_logging(stdio_level=logging.INFO,
17 | file_level=logging.DEBUG,
18 | filename=".easy.log",
19 | filepath=os.path.join(my_path, "logs")):
20 | os.makedirs(filepath, exist_ok=True)
21 | filename = os.path.join(filepath, filename)
22 | global _configured
23 | if _configured:
24 | warning("Re-configuring logging")
25 | stdio_handler = logging.StreamHandler()
26 | stdio_handler.setLevel(stdio_level)
27 | file_hanlder = logging.FileHandler(filename)
28 | file_hanlder.setLevel(file_level)
29 |
30 | logging.basicConfig(
31 | format="%(asctime)s - %(levelname)s - %(name)s - %(message).200s",
32 | datefmt="%m/%d/%Y %H:%M:%S",
33 | level=min(stdio_level, file_level),
34 | handlers=[stdio_handler, file_hanlder]
35 | )
36 |
37 | _configured = True
38 | _get_or_create_logger().debug("Configured logging")
39 |
40 |
41 | _loggers = {}
42 |
43 |
44 | def _get_or_create_logger():
45 | global _configured, _loggers
46 | if not _configured:
47 | configure_logging()
48 | try:
49 | for frame in inspect.stack():
50 | name = inspect.getmodule(frame[0]).__name__
51 | if name != __name__:
52 | break
53 | except:
54 | name = "_"
55 | if name not in _loggers:
56 | _loggers[name] = logging.getLogger(name)
57 | return _loggers[name]
58 |
59 |
60 | def print_to_string(*args, end="", **kwargs):
61 | with io.StringIO() as buf:
62 | print(*args, file=buf, end=end, **kwargs)
63 | return buf.getvalue()
64 |
65 |
66 | def debug(*args, **kwargs):
67 | _get_or_create_logger().debug(print_to_string(*args, **kwargs))
68 |
69 |
70 | def info(*args, **kwargs):
71 | _get_or_create_logger().info(print_to_string(*args, **kwargs))
72 |
73 |
74 | log = info
75 |
76 |
77 | def warning(*args, **kwargs):
78 | _get_or_create_logger().warning(print_to_string(*args, **kwargs))
79 |
80 |
81 | warn = warning
82 |
83 |
84 | def error(*args, **kwargs):
85 | _get_or_create_logger().error(print_to_string(*args, **kwargs))
86 |
--------------------------------------------------------------------------------
/solvers/gpt3/lm_solve/__init__.py:
--------------------------------------------------------------------------------
1 | from lm_solve.run import *
2 |
--------------------------------------------------------------------------------
/solvers/gpt3/lm_solve/gpt3_lib.py:
--------------------------------------------------------------------------------
1 | import os
2 | import json
3 | import openai
4 | import ezlog
5 | import time
6 | import datetime
7 |
8 | assert 'OPENAI_API_KEY' in os.environ, "Need to set environment variable `OPENAI_API_KEY`"
9 | openai.api_key = os.environ['OPENAI_API_KEY']
10 |
11 |
12 | _CACHE_PATH = os.path.join(os.path.dirname(__file__), "../.cache")
13 | _CACHE_FILENAME = os.path.join(_CACHE_PATH, "gpt3.cache")
14 | _ENCODING = "utf-8"
15 |
16 | _cache = None
17 |
18 |
19 | # the cache file is just a list of (query params dictionary encoded as a string but without n, result list)
20 | # multiple queries with the same params (except for n) are merged into a single big list
21 | def _save_line(item, comment=None):
22 | global _cache
23 | assert _cache is not None
24 | with open(_CACHE_FILENAME, "a", encoding=_ENCODING) as f:
25 | f.write(str(item)+ ((" # " + comment + "\n") if comment else "\n"))
26 |
27 | def _load_cache():
28 | global _cache
29 |
30 | assert _cache is None, "gpt3 cache already loaded"
31 |
32 | if not os.path.exists(_CACHE_PATH):
33 | ezlog.warn("Creating cache path")
34 | os.makedirs(_CACHE_PATH)
35 |
36 | _cache = {}
37 |
38 | if os.path.exists(_CACHE_FILENAME):
39 | time0 = time.perf_counter()
40 | with open(_CACHE_FILENAME, "r", encoding=_ENCODING) as f:
41 | for k, v in [eval(line) for line in f.readlines()]:
42 | if k not in _cache:
43 | _cache[k] = v
44 | else:
45 | _cache[k].extend(v)
46 | ezlog.info(f"Loaded gpt3 cache in {time.perf_counter()-time0:.1f}s")
47 | else:
48 | ezlog.warn("No gpt3 cache yet")
49 |
50 |
51 |
52 | def query(prompt, n=10, max_tokens=150, temp=1.0, max_batch=32, stop=None, notes=None, cache_only=False, verbose=True):
53 | """Query gpt3
54 |
55 | :param prompt: Up to 2048 tokens (about 3-4k chars)
56 | :param n: number of answers, None returns all cached answers
57 | :param max_tokens:
58 | :param temp: 0.9 seems to work well
59 | :param max_batch: max to query at once
60 | :param stop: string to stop at or '' if not to stop
61 | :param notes: notes you want to save or change in case you want to run the same query more than once!
62 | :return: list of answers and then the response items
63 | """
64 | global _cache
65 | if _cache is None:
66 | _load_cache()
67 |
68 | if temp == 0 and n > 1:
69 | ezlog.debug("Temp 0: no point in running more than one query")
70 | n = 1
71 |
72 | key = str(dict(prompt=prompt, max_tokens=max_tokens, temp=temp, max_batch=max_batch, stop=stop, rep=notes))
73 | cached = _cache.get(key, [])
74 | if n is None:
75 | return cached[:]
76 |
77 | if len(cached) >= n:
78 | return cached[:n]
79 |
80 | if cache_only:
81 | pass
82 | 1/0
83 | assert not cache_only, "Entry not found in cache"
84 | if verbose:
85 | print("/"*100)
86 | print("Querying GPT3 with prompt:")
87 | print(prompt)
88 | s = stop and stop.replace('\n', '\\n')
89 | print(f"/// n={n} ({n-len(cached)} new) max_tokens={max_tokens} temp={temp} max_batch={max_batch} stop={s}")
90 | print("/"*100)
91 |
92 | time0 = time.perf_counter()
93 |
94 | new = []
95 | n -= len(cached)
96 |
97 | while n > 0:
98 | m = min(n, max_batch)
99 |
100 | res = openai.Completion.create(
101 | engine="davinci-msft",
102 | prompt=prompt,
103 | max_tokens=max_tokens,
104 | temperature=temp,
105 | n=m,
106 | stop=stop or None
107 | )
108 |
109 | new += [c["text"] for c in res["choices"]]
110 | n -= m
111 |
112 | _save_line((key, new), f"{time.perf_counter() - time0:.1f}s {datetime.datetime.now()}")
113 | ans = _cache[key] = cached + new
114 | return ans[:]
115 |
116 | # old code
117 | # # to persist calls to the API...
118 | # _disk_cache = joblib.Memory(os.path.join(os.path.dirname(__file__), ".cache"), verbose=1).cache
119 | #
120 | #
121 | # @_disk_cache
122 | # def query(prompt, n=10, max_tokens=150, temperature=1.0, max_batch=32):
123 | # """Query gpt3
124 | #
125 | # :param prompt: Up to 2048 tokens (about 3-4k chars)
126 | # :param n: number of answers
127 | # :param max_tokens:
128 | # :param temperature:
129 | # :param max_batch: max to query at once
130 | # :return: list of answers and then the response items
131 | # """
132 | # if temperature == 0 and n > 1:
133 | # ezlog.debug("Temp 0: no point in running more than one query")
134 | # n = 1
135 | #
136 | # responses = []
137 | # while n > 0:
138 | # m = min(n, max_batch)
139 | # prompt_summary = prompt if len(prompt) < 80 else f"{prompt[:40]}...{prompt[-40:]}"
140 | # ezlog.warn(f"**** Running GPT3 query: temp {temperature}, n={m}, prompt={prompt_summary}")
141 | # time0 = time.perf_counter()
142 | # responses.append(openai.Completion.create(
143 | # engine="davinci-msft",
144 | # prompt=prompt,
145 | # max_tokens=max_tokens,
146 | # temperature=temperature,
147 | # n=m
148 | # ))
149 | # ezlog.info(f"**** Got response in {time.perf_counter()-time0}s...")
150 | # n -= m
151 | #
152 | # return [c["text"] for r in responses for c in r["choices"]], responses
153 |
154 |
155 |
--------------------------------------------------------------------------------
/solvers/gpt3/lm_solve/judge.py:
--------------------------------------------------------------------------------
1 | import hashlib
2 | import os
3 | import time
4 | from pebble import ProcessPool
5 | import multiprocessing
6 | from concurrent.futures import TimeoutError
7 | from typing import List, Set, Tuple, Dict, Union, Any
8 |
9 | import utils
10 | import ezlog
11 |
12 | def _COPY(x):
13 | return x
14 |
15 | _ENV = dict(List=List, Set=Set, Tuple=Tuple, Dict=Dict, COPY=_COPY,
16 | os=None, sys=None, input=None, open=None, print=None, compile=None, copyright=None)
17 | _UNSAFE = ["import", "builtin", "__class"]
18 |
19 | MAX_WORKERS = multiprocessing.cpu_count() // 2
20 |
21 | _CACHE_PATH = os.path.join(os.path.dirname(__file__), "../.cache") # ".cache"
22 |
23 |
24 | class SetCache:
25 | """Simple cache that stores a set of keys. Cannot remove values. Haven't yet implemented iteration."""
26 | BYTE_LEN = 256 // 8 # for sha256
27 |
28 | def __init__(self, name, path=_CACHE_PATH):
29 | self._path = path
30 | self.name = name
31 | self._filename = os.path.join(path, f"{name}_set.cache")
32 | self._set = None # the main set, loaded lazily
33 |
34 | def update(self, keys):
35 | self._load()
36 | hashes = [self._hash(k) for k in keys]
37 | additions = {h for h in hashes if h not in self._set}
38 |
39 | if additions:
40 | with open(self._filename, "ab") as f: # append to binary file
41 | for h in additions:
42 | f.write(h)
43 | self._set.update(additions)
44 |
45 | def add(self, key):
46 | self.update([key])
47 |
48 | def _hash(self, key):
49 | return hashlib.sha256(bytes(str(key), encoding='utf8')).digest()
50 |
51 | def __contains__(self, key: str):
52 | self._load()
53 | h = self._hash(key)
54 | return h in self._set
55 |
56 | def __delitem__(self, key):
57 | raise NotImplementedError
58 |
59 | def __len__(self):
60 | self._load()
61 | return len(self._set)
62 |
63 | def _load(self):
64 | if self._set is None: # only load if not already loaded
65 | if not os.path.exists(self._path):
66 | ezlog.warn(f"Creating path for `{self.name}` cache")
67 | os.makedirs(self._path)
68 |
69 | time0 = time.perf_counter()
70 | if os.path.exists(self._filename):
71 | with open(self._filename, "rb") as f: # read binary file
72 | data = f.read()
73 | self._set = {data[j:j + self.BYTE_LEN] for j in range(0, len(data), self.BYTE_LEN)}
74 | else:
75 | self._set = set()
76 | dur = time.perf_counter() - time0
77 | ezlog.info(f"Loaded `{self.name}` cache of {len(self):,} items in {dur:.1f}s")
78 |
79 |
80 | def _judge(code_env):
81 | code, env = code_env
82 | for u in _UNSAFE:
83 | if u in code:
84 | return False
85 |
86 | try:
87 | exec(code, env.copy()) # not sure if copy() is necessary
88 | return True
89 | except Exception as e:
90 | return False
91 |
92 |
93 | # Cache judge results (which are nondeterministic due to timeout) for reproducibility
94 | _judge_success = SetCache('judge_success')
95 | _judged_batches = SetCache('judged_batches')
96 |
97 | def judge_parallel(src_codes, timeout, max_workers=MAX_WORKERS, env=_ENV, force_compute=False):
98 | global _judge_success, _judged_batches
99 |
100 | if force_compute or src_codes not in _judged_batches:
101 | new_codes = utils.dedup(code for code in src_codes if code not in _judge_success)
102 | if new_codes:
103 | ezlog.info(f"Judging {len(new_codes)}/{len(src_codes)} new codes (removing duplicates/things in cache)")
104 | successes = []
105 |
106 |
107 | with ProcessPool(max_workers=max_workers) as pool:
108 | future = pool.map(_judge, [(code, env) for code in new_codes], timeout=timeout)
109 |
110 | results = future.result()
111 |
112 | i = 0
113 | while True:
114 | try:
115 | if next(results):
116 | successes.append(new_codes[i])
117 | except StopIteration:
118 | _judge_success.update(successes)
119 | break
120 | except (TimeoutError, Exception) as error:
121 | pass
122 | assert i < len(new_codes)
123 | i += 1
124 | assert i == len(new_codes)
125 | _judged_batches.add(src_codes)
126 |
127 | return [code in _judge_success for code in src_codes]
128 |
129 |
130 | if __name__ == "__main__":
131 | res = judge_parallel([
132 | """1+1
133 | """,
134 | """assert False,'cats'""",
135 | """assert False""",
136 | """1[2]""",
137 | """1/0""",
138 | """while True:
139 | pass""",
140 | """for i in range(10**5):
141 | pass"""
142 | ], timeout=1.0)
143 | print(res)
144 |
--------------------------------------------------------------------------------
/solvers/gpt3/requirements.txt:
--------------------------------------------------------------------------------
1 | astor==0.8.1
2 | numpy==1.22.0
3 | openai==0.6.3
4 | tqdm==4.60.0
5 | transformers==4.30.0
6 | Pebble==4.6.1
7 |
8 | # we ran with Python version sys.version = '3.6.9 (default, Jan 26 2021, 15:33:00) \n[GCC 8.4.0]'
9 | # distro-info===0.18ubuntu0.18.04.1
10 |
--------------------------------------------------------------------------------
/solvers/gpt3/run_gpt3_experiments.py:
--------------------------------------------------------------------------------
1 | """
2 | This script runs the GPT_3 experiments and prints the results to stdout.
3 | It uses cacheing mechanisms so that if run twice with the same parameters, it will give exactly the same
4 | results and will not query the GPT3 API again and will not judge the resulting solutions again. Hence, the first
5 | time you run it, it will be slow, but you can subsequently run it again and it will be fast. It will run the
6 | experiment three times, with different seeds to get different results.
7 |
8 | """
9 |
10 | import lm_solve
11 | import numpy as np
12 |
13 | PARAMS = dict(
14 | temp=0.9,
15 | timeout=1.0, # seconds to judge
16 | n=10 * 1000,
17 | filename="puzzles_with_descriptions.json",
18 | stop="\n",
19 | cache_only=False, # change this to True if you want to run a 2nd time without risking hitting API
20 | )
21 |
22 | BOOTSTRAP_PARAMS = dict(
23 | temp=PARAMS["temp"],
24 | timeout=PARAMS["timeout"],
25 | filename=PARAMS["filename"],
26 | stop=PARAMS["stop"],
27 | cache_only=PARAMS["cache_only"],
28 | ppi=32, # puzzles per iteration
29 | iterations=(PARAMS["n"] + 31) // 32,
30 | )
31 |
32 | STUDY = range(107, 137) # the range of puzzles used in the study
33 |
34 | PREFIX = """
35 | def f1(s: str):
36 | return "Hello " + s == "Hello world"
37 |
38 | assert True == f1("world")
39 |
40 | ---
41 |
42 | def f2(s: str):
43 | return "Hello " + s[::-1] == "Hello world"
44 |
45 | assert True == f2("world"[::-1])
46 |
47 | ---
48 |
49 | def f3(x: List[int]):
50 | return len(x) == 2 and sum(x) == 3
51 |
52 | assert True == f3([1, 2])
53 |
54 | ---
55 |
56 | def f4(s: List[str]):
57 | return len(set(s)) == 1000 and all((x.count("a") > x.count("b")) and ('b' in x) for x in s)
58 |
59 | assert True == f4(["a"*(i+2)+"b" for i in range(1000)])
60 |
61 | ---
62 |
63 | def f5(n: int):
64 | return str(n * n).startswith("123456789")
65 |
66 | assert True == f5(int(int("123456789" + "0"*9) ** 0.5) + 1)
67 |
68 | ---
69 |
70 | """ # trailing newlines important
71 |
72 | PREFIX_DOCSTR = '''
73 | def f1(s: str):
74 | """Find a string that when concatenated onto 'Hello ' gives 'Hello world'."""
75 | return "Hello " + s == "Hello world"
76 |
77 | assert True == f1("world")
78 |
79 | ---
80 |
81 | def f2(s: str):
82 | """Find a string that when reversed and concatenated onto 'Hello ' gives 'Hello world'."""
83 | return "Hello " + s[::-1] == "Hello world"
84 |
85 | assert True == f2("world"[::-1])
86 |
87 | ---
88 |
89 | def f3(x: List[int]):
90 | """Find a list of two integers whose sum is 3."""
91 | return len(x) == 2 and sum(x) == 3
92 |
93 | assert True == f3([1, 2])
94 |
95 | ---
96 |
97 | def f4(s: List[str]):
98 | """Find a list of 1000 distinct strings which each have more 'a's than 'b's and at least one 'b'."""
99 | return len(set(s)) == 1000 and all((x.count("a") > x.count("b")) and ('b' in x) for x in s)
100 |
101 | assert True == f4(["a"*(i+2)+"b" for i in range(1000)])
102 |
103 | ---
104 |
105 | def f5(n: int):
106 | """Find an integer whose perfect square begins with 123456789 in its decimal representation."""
107 | return str(n * n).startswith("123456789")
108 |
109 | assert True == f5(int(int("123456789" + "0"*9) ** 0.5) + 1)
110 |
111 | ---
112 |
113 | ''' # trailing newlines important
114 |
115 |
116 | def run(seed=0):
117 | sols = [lm_solve.prompt_experiment(**PARAMS, prefix="", seed=seed),
118 | lm_solve.prompt_experiment(**PARAMS, prefix=PREFIX, seed=seed),
119 | lm_solve.prompt_experiment(**PARAMS, prefix=PREFIX_DOCSTR, add_docstring=True, seed=seed)]
120 | problems_solved = [sorted([i for i, (f, gs) in enumerate(s) if gs]) for s in sols]
121 | bootstrap = lm_solve.bootstrap(**BOOTSTRAP_PARAMS, seed=seed)
122 | print(f"run={seed} ALL DONE!\n\n")
123 | print(f"run={seed} RESULTS " + "=" * 50)
124 | print()
125 |
126 | # Instead of running until first success, and outputting number of attempts, we do something more accurate.
127 | # We run for N tries for each problem and do not stop on first success. Then, we use the number of successes
128 | # for a better estimate of the average number of attempts required for first success. If we have s successes
129 | # out of N attempts, then the expected number of attempts is (N - s) / (1 + s). This is the expectation of the
130 | # random variable that is: when you permute the attempts uniformly at random, how many attempts before the
131 | # first success. If s=N, it's 0, if s=1, it's (N-1)/2, etc.
132 | counts = [[(PARAMS["n"] - len(gs)) / (1 + len(gs)) for f, gs in s if gs] for s in sols]
133 | counts.append([m for m, _i, _f, _a in bootstrap])
134 | counts = [[1 + z for z in c] for c in counts] # add 1 to make it 1-based
135 | for c in counts:
136 | c.sort()
137 | print(f"run={seed} (Expected) number of attempts before a problem is solved [short, med, long, bootstrap]:")
138 | print(counts)
139 | problems_solved.append([i for _m, i, _f, _a in bootstrap])
140 | print()
141 | print(f"run={seed} Which problems were solved [short, med, long, bootstrap]:")
142 | print(problems_solved)
143 | print()
144 | print(f"run={seed} Number of problems solved [short, med, long, bootstrap]:")
145 | print([len(c) for c in counts])
146 | print()
147 | print(f"run={seed} Number of 30 study problems solved [short, med, long, bootstrap]:")
148 | print([len([i for i in s if i in STUDY]) for s in problems_solved])
149 | print()
150 | difficulties = [1.0 for _ in range(len(sols[0]))]
151 |
152 | k = 1
153 | for m, i, f, a in bootstrap:
154 | difficulties[i] = np.log(m + 1) / np.log(PARAMS["n"])
155 |
156 | # These commented lines print the problems that bootstrap solved
157 | # print()
158 | # print(f"# Bootstrap solved after {m + 1} tries:")
159 | # print(f.replace("def f", "def f" + str(k)))
160 | # import json
161 | #
162 | # print(f"SOL:", json.dumps(a))
163 | k += 1
164 | print(f"run={seed} Bootstrap difficulties for study puzzles:")
165 | print([difficulties[i] for i in STUDY])
166 |
167 |
168 | if __name__ == "__main__":
169 | for seed in range(3):
170 | run(seed)
171 |
--------------------------------------------------------------------------------
/solvers/gpt3/utils/__init__.py:
--------------------------------------------------------------------------------
1 | from utils.str_utils import *
2 | from utils.time_utils import *
3 |
4 | import functools
5 | import operator
6 |
7 |
8 | def prod(iterable): # like sum but product
9 | return functools.reduce(operator.mul, iterable, 1)
10 |
11 |
12 | def flatten(it):
13 | return (e for a in it for e in (flatten(a) if isinstance(a, (tuple, list)) else (a,)))
14 |
15 |
16 | def load_json(filename):
17 | import json
18 | with open(filename, "r") as f:
19 | return json.load(f)
20 |
21 |
22 | def viz_py(py):
23 | import astor, ast
24 | print(astor.dump_tree(ast.parse(py)))
25 |
26 |
27 | def dedup(li):
28 | seen = set()
29 | return [x for x in li if x not in seen and not seen.add(x)]
30 |
--------------------------------------------------------------------------------
/solvers/gpt3/utils/str_utils.py:
--------------------------------------------------------------------------------
1 | import json
2 |
3 | def get_lambda_arg_name(lam):
4 | assert lam.startswith("lambda ")
5 | return lam[len("lambda "):lam.index(":")].strip()
6 |
7 |
8 | def stringify(const):
9 | if type(const) is str:
10 | return json.dumps(const)
11 | return str(const)
12 |
13 |
14 | def color_str(obj, code="\033[0;36m"):
15 | return code + str(obj) + '\033[0m'
16 |
17 |
--------------------------------------------------------------------------------
/solvers/gpt3/utils/time_utils.py:
--------------------------------------------------------------------------------
1 | import time
2 | import logging
3 |
4 | logger = logging.getLogger(__name__)
5 |
6 | def timeit(method):
7 | def timed(*args, **kw):
8 | tick = time.time()
9 | result = method(*args, **kw)
10 | tock = time.time()
11 | logger.debug(f'{method.__name__}: {tock - tick:.3f}s')
12 |
13 | return result
14 | return timed
15 |
--------------------------------------------------------------------------------
/utils.py:
--------------------------------------------------------------------------------
1 | import json
2 | import logging
3 | import inspect
4 | import io
5 | import os
6 |
7 | my_path = os.path.dirname(__file__)
8 |
9 | def check_hashseed(desired_seed = 0):
10 | if os.environ.get('PYTHONHASHSEED') != desired_seed:
11 | info(f"Ideally set PYTHONHASHSEED={desired_seed} for perfect reproducibility")
12 | return False
13 | return True
14 |
15 |
16 | def inv_dict(d):
17 | ans = {}
18 | for k, v in d.items():
19 | if v not in ans:
20 | ans[v] = []
21 | ans[v].append(k)
22 | return ans
23 |
24 |
25 | def remove_docstring(f):
26 | """Remove docstring"""
27 | assert '\n """' in f, f"No triple quote docstring (after four spaces) in: \n{f}"
28 | i = f.index('\n """')
29 | j = f.index('"""', i + 8)
30 | return f[:i + 1] + f[j + 4:]
31 |
32 |
33 | def get_docstring(f):
34 | assert '\n """' in f, f"No triple quote docstring (after four spaces) in: \n{f}"
35 | i = f.index('\n """')
36 | j = f.index('"""', i + 8)
37 | docstring = f[i + 1:j + 3]
38 | if not docstring.strip(' "'):
39 | warn(f"Empty docstring in:\n{f}")
40 | return docstring
41 |
42 |
43 | def flatten(it):
44 | return (e for a in it for e in (flatten(a) if isinstance(a, (tuple, list)) else (a,)))
45 |
46 | def save_json(obj, filename, make_dirs_if_necessary=False, indent=2, **kwargs):
47 | """Saves compressed file if filename ends with '.gz'"""
48 | import json
49 | if make_dirs_if_necessary:
50 | os.makedirs(os.path.dirname(filename), exist_ok=True)
51 | if filename.endswith(".gz"):
52 | import gzip
53 | with gzip.open(filename, "wt") as f:
54 | return json.dump(obj, f, indent=indent, **kwargs)
55 | with open(filename, "w", encoding="utf8") as f:
56 | return json.dump(obj, f, indent=indent, **kwargs)
57 |
58 | def load_json(filename):
59 | """Loads compressed file if filename ends with '.gz'"""
60 | import json
61 | if filename.endswith(".gz"):
62 | import gzip
63 | with gzip.open(filename, "rt") as f:
64 | return json.load(f)
65 | with open(filename, "r", encoding="utf8") as f:
66 | return json.load(f)
67 |
68 |
69 | def stringify(const):
70 | if type(const) is str:
71 | return json.dumps(const)
72 | return str(const)
73 |
74 |
75 | def dedup(stuff):
76 | seen = set()
77 | return [a for a in stuff if a not in seen and not seen.add(a)]
78 |
79 |
80 | def color_str(obj, code="\033[0;36m"):
81 | return code + str(obj) + '\033[0m'
82 |
83 |
84 | _configured = False
85 |
86 |
87 | def configure_logging(stdio_level=logging.INFO,
88 | file_level=logging.DEBUG,
89 | filename=os.path.join(my_path, ".problems.log")):
90 | global _configured
91 | if _configured:
92 | warning("Re-configuring logging")
93 | stdio_handler = logging.StreamHandler()
94 | stdio_handler.setLevel(stdio_level)
95 | file_hanlder = logging.FileHandler(filename)
96 | file_hanlder.setLevel(file_level)
97 |
98 | logging.basicConfig(
99 | format="%(asctime)s - %(levelname)s - %(name)s - %(message).200s",
100 | datefmt="%m/%d/%Y %H:%M:%S",
101 | level=min(stdio_level, file_level),
102 | handlers=[stdio_handler, file_hanlder]
103 | )
104 |
105 | _configured = True
106 | _get_or_create_logger().debug("Configured logging")
107 |
108 |
109 | _loggers = {}
110 |
111 |
112 | def _get_or_create_logger():
113 | global _configured, _loggers
114 | if not _configured:
115 | configure_logging()
116 | name = "_"
117 | for frame in inspect.stack():
118 | name = inspect.getmodule(frame[0]).__name__
119 | if name != __name__:
120 | break
121 | if name not in _loggers:
122 | _loggers[name] = logging.getLogger(name)
123 | return _loggers[name]
124 |
125 |
126 | def print_to_string(*args, end="", **kwargs):
127 | with io.StringIO() as buf:
128 | print(*args, file=buf, end=end, **kwargs)
129 | return buf.getvalue()
130 |
131 |
132 | def debug(*args, **kwargs):
133 | _get_or_create_logger().debug(print_to_string(*args, **kwargs))
134 |
135 |
136 | def info(*args, **kwargs):
137 | _get_or_create_logger().info(print_to_string(*args, **kwargs))
138 |
139 |
140 | log = info
141 |
142 |
143 | def warning(*args, **kwargs):
144 | _get_or_create_logger().warning(print_to_string(*args, **kwargs))
145 |
146 |
147 | warn = warning
148 |
149 |
150 | def error(*args, **kwargs):
151 | _get_or_create_logger().error(print_to_string(*args, **kwargs))
152 |
--------------------------------------------------------------------------------