Duplicate from codeparrot/apps_metric

Co-authored-by: Loubna Ben Allal <[email protected]>

.gitattributes

@@ -0,0 +1,27 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,53 @@
+---
+title: APPS Metric
+emoji: 📊
+colorFrom: blue
+colorTo: pink
+datasets:
+- null
+tags:
+- evaluate
+- metric
+description: Evaluation metric for the APPS benchmark
+sdk: gradio
+sdk_version: 3.0.2
+app_file: app.py
+pinned: false
+duplicated_from: codeparrot/apps_metric
+---
+
+# Metric Card for apps_metric 
+
+## Metric Description
+This metric is used to evaluate code generation on the [APPS benchmark](https://huggingface.co/datasets/codeparrot/apps).
+
+## How to Use
+You can load the metric and use it with the following commands:
+
+```python
+from evaluate import load
+apps_metric = load('codeparrot/apps_metric')
+# to evaluate generations made for all levels for example
+results = apps_metric.compute(predictions=generations, level="all")
+```
+
+### Inputs
+**generations** list(list(str)): List of code generations, each sub-list corresponds to the generations for a problem in APPS dataset, **the order of the samples in the dataset must be kept (with respect to the difficulty level)**.
+
+### Output Values
+
+**average accuracy**: when a single solution is generated, average accuracy computes the average of test cases that are passed.
+
+**strict accuracy**: when a single solution is generated, strict accuracy computes the average number of problems that pass all their test cases.
+
+**pass@k**: when multiple solutions are generated per problem, pass@k is the metric originally used for the [HumanEval](https://huggingface.co/datasets/openai_humaneval) benchmark. For more details please refer to the [metric space](https://huggingface.co/spaces/evaluate-metric/code_eval) and [Codex paper](https://arxiv.org/pdf/2107.03374v2.pdf).
+
+## Citation
+```
+@article{hendrycksapps2021,
+  title={Measuring Coding Challenge Competence With APPS},
+  author={Dan Hendrycks and Steven Basart and Saurav Kadavath and Mantas Mazeika and Akul Arora and Ethan Guo and Collin Burns and Samir Puranik and Horace He and Dawn Song and Jacob Steinhardt},
+  journal={NeurIPS},
+  year={2021}
+}
+``` 

app.py

@@ -0,0 +1,6 @@
+import evaluate
+from evaluate.utils import launch_gradio_widget
+
+
+module = evaluate.load("loubnabnl/apps_metric")
+launch_gradio_widget(module)

apps_metric.py

@@ -0,0 +1,82 @@
+# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Evaluation of code generation on the APPS benchmark"""
+
+import evaluate
+import datasets
+from .utils import compute_metrics
+from .testing_util import run_test
+
+
+_CITATION = """\
+@article{hendrycksapps2021,
+  title={Measuring Coding Challenge Competence With APPS},
+  author={Dan Hendrycks and Steven Basart and Saurav Kadavath and Mantas Mazeika and Akul Arora and Ethan Guo and Collin Burns and Samir Puranik and Horace He and Dawn Song and Jacob Steinhardt},
+  journal={NeurIPS},
+  year={2021}
+}
+"""
+
+
+_DESCRIPTION = """\
+This is a metric to evaluate code generation using the APPS benchmark "Measuring Coding Challenge Competence With
+APPS" (https://arxiv.org/pdf/2105.09938.pdf).
+"""
+
+
+# TODO: Add description of the arguments of the module here
+_KWARGS_DESCRIPTION = """
+Computes Average accuracy and strict accuracy for single generations, and pass@k for multiple generations.
+Args:
+    predictions: list of code generations to score. It's a list of list(s), each corresponding to a problem from APPS dataset.
+
+Returns:
+    metrics: dict of three metrics: average accuracy, stric accuracy, and pass@k.
+Examples:
+    >>> my_new_module = evaluate.load("loubnabnl/apps_metric")
+    >>> results = my_new_module.compute(predictions=[["s=input()\nprint(s)"]])
+    >>> print(results)
+    {'avg_accuracy': 0, 'strict_accuracy': 0, 'pass_at_k': None}
+"""
+
+
+
+
+@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
+class apps_metric(evaluate.EvaluationModule):
+    """Evaluate code generation on APPS benchmark. 
+    The generations are compiled and their corresponding unit tests are run"""
+
+    def _info(self):
+
+        return evaluate.EvaluationModuleInfo(
+
+            module_type="metric",
+            description=_DESCRIPTION,
+            citation=_CITATION,
+            inputs_description=_KWARGS_DESCRIPTION,
+
+            features=datasets.Features({
+                'predictions': datasets.Sequence(datasets.Value("string")),
+            }),
+            homepage="https://github.com/hendrycks/apps",
+            reference_urls=["https://huggingface.co/datasets/codeparrot/apps"]
+        )
+
+
+
+    def _compute(self, predictions, k_list=[1, 10, 100], count_errors=True, level="all", debug=False):
+        """Returns the scores"""
+        metrics = compute_metrics(predictions, k_list=k_list, count_errors=count_errors, level=level, debug=debug)
+        return metrics

example_script.py

@@ -0,0 +1,133 @@
+"""This is an example script to evaluate a code generation model on APPS, you can also use the APPS solutions as code generations
+> python example_script.py --model_ckpt MODEL_NAME --num_tasks 10 --difficulty introductory --n_samples 1
+> python example_script.py --use_solutions True --num_tasks 10 --difficulty introductory --n_samples 1"""
+
+import json
+import pprint
+from tqdm import tqdm
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, set_seed
+from evaluate import load
+
+def generate_prompt(sample):
+    starter_code = None if len(sample["starter_code"]) == 0 else sample["starter_code"] 
+    try:
+        input_outpout = json.loads(sample["input_output"])
+        fn_name = None if not input_outpout.get("fn_name") else input_outpout["fn_name"] 
+    except ValueError:
+        fn_name = None 
+    _input = "\nQUESTION:\n"
+    _input += sample["question"]
+    if starter_code:
+        _input += starter_code
+    if fn_name:
+        _input += "\nUse Standard Input format"
+    else:
+        _input += "\nUse Call-Based format"
+    
+    _input += "\nANSWER:\n"
+    return _input
+
+
+def complete_code(pipe, prompt, num_completions=1, max_length=256, **gen_kwargs):
+    """Complete prompt with text generation pipeline and return num_completions."""
+    prompt = pipe.tokenizer.eos_token + prompt
+    try:
+        code_gens = pipe(prompt, num_return_sequences=num_completions, max_length=max_length, **gen_kwargs)
+        return [code_gen["generated_text"][len(prompt):] for code_gen in code_gens]
+    except IndexError:
+        print("prompt is longer than the context size of the model, generation skipped")
+        code_gens = ""
+        return [""]
+
+
+def make_generations(dataset, args, model, tokenizer):
+    set_seed(args.seed)
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=args.device_int)
+
+    # Generation settings
+    gen_kwargs = {
+        "do_sample": args.do_sample,
+        "temperature": args.temperature,
+        "top_p": args.top_p,
+        "top_k": args.top_k
+    }
+
+    # Generate completions for evaluation set
+    n_tasks = args.num_tasks if args.num_tasks is not None else len(dataset)
+    print(f"ntasks is {n_tasks}")
+    generations = []
+    for task in tqdm(range(n_tasks)):
+        task_generations = []
+        prompt = generate_prompt(dataset[task]).strip()
+        task_generations.extend(complete_code(pipe, prompt, num_completions=args.n_samples, max_length=args.max_length, **gen_kwargs))
+        generations.append([gen.replace(args.eos, "") for gen in task_generations])
+    return generations
+
+
+def main(args):
+    DATA_PATH = "codeparrot/apps"
+    argsdict = vars(args)
+    print(pprint.pformat(argsdict))
+
+    # setup
+    print("Loading evaluation dataset...")
+    dataset = load_dataset(DATA_PATH, split="test", difficulties=[args.difficulty])
+    if args.use_solutions:
+        print("Using data solutions as code generations")
+        model = None
+        tokenizer = None
+        generations = []
+        for index in range(args.num_tasks+1):
+            try:
+                sol = json.loads(dataset[index]["solutions"])
+                generations.append(sol[:args.n_solutions])
+            except ValueError:
+                print(f"No solutions for task {index} or not enough to have {args.n_solutions} solutions")
+                break
+        
+    else:
+        print("Loading tokenizer and model...")
+        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
+        model = AutoModelForCausalLM.from_pretrained(args.model_ckpt)
+        generations = make_generations(dataset, args, model, tokenizer)
+        
+    metric = load("loubnabnl/apps_metric")
+    results = metric.compute(predictions=generations, level=args.difficulty, k_list=args.k_list, count_errors=args.count_errors, debug=args.debug)
+    print(results)
+    with open(args.output_file, "w") as fp:
+        json.dump(results, fp)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Testing a Language Model on APPS Python Code dataset")
+    #model and tokenizer arguments
+    parser.add_argument("--model_ckpt", default="loubnabnl/apps-1.5B-model", type=str, help="path to model checkpoint.")
+    parser.add_argument("--tokenizer", default="gpt2", type=str, help="tokenizer to use.")
+    parser.add_argument("--eos", default="<|endoftext|>", type=str, help="end of sentence token.")
+    # generation arguments
+    parser.add_argument("--do_sample", default=True, type=bool, help="do sampling in generation")
+    parser.add_argument("--temperature", default=0.2, type=float, help="temperature for sampling")
+    parser.add_argument("--top_p", default=0.95, type=float, help="top p for sampling")
+    parser.add_argument("--top_k", default=0, type=float, help="top k for sampling")
+    parser.add_argument("--max_length", default=1024, type=int, help="max length of generated code")
+    # evaluation arguments
+    parser.add_argument("--difficulty", default="all", type=str, help="difficulty level to select in the dataset from:\
+     'all', 'introductory', 'interview'  and 'competition' ")
+    parser.add_argument("--num_tasks", default=6, type=int, help="number of tasks to evaluate")
+    parser.add_argument("--use_solutions", default=False, type=bool, help="use solutions instead of generating new code")
+    parser.add_argument("--n_samples", default=1, type=int, help="number of samples to generate")
+    parser.add_argument("--n_solutions", default=1, type=int, help="number of solutions to use")
+    parser.add_argument("--k_list", default=[1, 2, 3], type=list, help="list of k values to evaluate pass@k")
+    parser.add_argument("--count_errors", default=False, type=bool, help="count compilation and runtime errors for single generations")
+    # configuration
+    parser.add_argument("--seed", default=0, type=int, help="generation seed")
+    parser.add_argument("--device_int", default=-1, type=int, help="device on which code generation is run, if positive use GPU")
+    parser.add_argument("--debug", default=False, type=bool, help="debug mode")
+    # save
+    parser.add_argument("--output_file", default="apps_metrics.json", type=str, help="output file to save the results")
+
+    args = parser.parse_args()
+    main(args)

requirements.txt

@@ -0,0 +1,3 @@
+evaluate==0.1.0
+datasets~=2.0
+pyext==0.7

test_examples/solutions_problem_1.json

@@ -0,0 +1 @@
+["s = input()\nn = len(s)\nind = -1\nf = False\nfor i in range(n):\n    if s[i] == '[':\n        f = True\n    elif s[i] == ':':\n        if f:\n            ind = i\n            break\nbind = -1\nf = False\nfor i in range(n-1,-1,-1):\n    if s[i] == ']':\n        f = True\n    elif s[i] == ':':\n        if f:\n            bind = i\n            break\n# print(ind,bind)\nif ind == -1 or bind == -1:\n    print(-1)\nelif ind >= bind:\n    print(-1)\nelse:\n    ans = 4\n    for i in range(ind+1,bind):\n        if s[i] == '|':\n            ans += 1\n    print(ans)\n", "def main():\n    s = input()\n    \n    if s.count('[') == 0 or s.count(']') == 0:\n        print(-1)\n        return\n    \n    t = s[s.find('['):s.rfind(']')+1]\n    \n    if t.count(':') < 2:\n        print(-1)\n        return\n    \n    t = t[t.find(':'):t.rfind(':')+1]\n    print(4 + t.count('|'))\n\nmain()", "s = input()\nif '[' in s:\n    s = s[s.find('[') + 1:]\n    if ']' in s:\n        s = s[:s.rfind(']')]\n        if s.count(':') >= 2:\n            s = s[s.find(':') + 1 : s.rfind(':')]\n            print(s.count('|') + 4)\n\n        else:\n            print(-1)\n    else:\n        print(-1)\nelse:\n    print(-1)", "import sys\ns = input()\nst = s.find('[')\nif st==-1: print((-1)); return\ns = s[st+1:]\n#print(s)\nst = s.find(':')\nif st==-1: print((-1)); return\ns = s[st+1:]\n#print(s)\ns = s[::-1]\nst = s.find(']')\nif st==-1: print((-1)); return\ns = s[st+1:]\n#print(s)\nst = s.find(':')\nif st==-1: print((-1)); return\ns = s[st+1:]\n#print(s)\nx = s.count('|')\nprint(x+4 if x>=0 else -1)\n", "s = input()\n\nsb,eb,sc,ec = -1, -1, -1, -1\n\nfor i in range(len(s)):\n\tif s[i] == '[' and sb == -1:\n\t\tsb = i\n\telif s[i] == ']':\n\t\teb = i\n\telif s[i] == ':' and sc == -1 and sb!=-1:\n\t\tsc = i\n\nif eb <= sb or sc>eb:\n\tprint(-1)\nelif sb ==-1 or eb==-1 or sc==-1:\n\tprint(-1)\nelse:\n\tfor i in range(sc+1, eb):\n\t\tif s[i] == ':':\n\t\t\tec = i\n\tif ec == -1:\n\t\tprint(-1)\n\telse:\n\t\tcnt = 0\n\t\tfor i in range(sc,ec):\n\t\t\tif (s[i] == '|'):\n\t\t\t\tcnt += 1\n\t\tprint(cnt+4)", "s = input()\nt_d = 0\ntry:\n    left = -1\n    was_b = False\n    for i in range(len(s)):\n        if s[i] == '[' and not was_b:\n            was_b = True\n            continue\n        if s[i] == ':' and was_b:\n            left = i\n            break\n        t_d += 1\n    if left == -1:\n        raise ArithmeticError()\n    right = -1\n    was_b = False\n    for i in range(len(s) - 1, -1, -1):\n        if s[i] == ']' and not was_b:\n            was_b = True\n            continue\n        if s[i] == ':' and was_b:\n            right = i\n            break\n        t_d += 1\n    if right == -1 or right <= left:\n        raise ArithmeticError()\n    for i in range(left + 1, right):\n        if s[i] != '|':\n            t_d += 1\n    print(len(s) - t_d)\nexcept:\n    print(-1)\n        \n", "s = input()\n\nmode = 0\nl = len(s)\nr = -1\nfor i in range(len(s)):\n    if mode == 0:\n        if s[i] == \"[\":\n            mode = 1\n    if mode == 1:\n        if s[i] == \":\":\n            l = i\n            break\n\nmode = 0\nfor i in range(len(s)-1, -1, -1):\n    if mode == 0:\n        if s[i] == \"]\":\n            mode = 1\n    if mode == 1:\n        if s[i] == \":\":\n            r = i\n            break\n    \nif l >= r:\n    print(-1)\nelse:\n    c = 0\n    for i in range(l+1, r):\n        if s[i] == \"|\":\n            c += 1\n    print(c+4)\n", "s = input()\n\nf1 = False\nf2 = False\nl1 = -1\nfor l in range(len(s)):\n    if f1 == False and s[l] == '[':\n        f1 = True\n    elif f1 == True and s[l] == ':':\n        f2 = True\n        l1 = l\n        break\ng1 = False\ng2 = False\nr1 = -1\nfor r in range(len(s) - 1, -1, -1):\n    if g1 == False and s[r] == ']':\n        g1 = True\n    elif g1 == True and s[r] == ':':\n        g2 = True\n        r1 = r\n        break\nif (l1 == -1 or r1 == -1) or (r1 <= l1):\n    print(-1)\n    \nelse:\n    ans = 4\n    for i in range(l1 + 1, r1):\n        if s[i] == '|': ans += 1\n    print(ans)", "s=input()\npos1=-1\npos2=-1\npos3=-1\npos4=-1\nfor i in range(0,len(s)):\n    if(s[i]=='['):\n        pos1=i\n        break\nfor i in range(len(s)-1,pos1,-1):\n    if(s[i]==']'):\n        pos2=i\n        break\nfor i in range(pos1,pos2+1):\n    if(s[i]==':'):\n        pos3=i\n        break\nfor i in range(pos2,pos3,-1):\n    if(s[i]==':'):\n        pos4=i\n        break\n    \nif(pos1==-1 or pos2==-1 or pos3==-1 or pos4==-1 or len(s)<4):\n    print('-1')\nelse:\n    c=0\n    for j in range(pos3,pos4):\n        if(s[j]=='|'):\n            c=c+1\n    print(c+4)\n", "def ii():\n    return int(input())\ndef mi():\n    return list(map(int, input().split()))\ndef li():\n    return list(mi())\n\ns = input().strip()\nn = len(s)\nans = -1\nfb = s.find('[')\nif fb >= 0:\n    fc = s.find(':', fb)\n    if fc >= 0:\n        lb = s.rfind(']')\n        if lb > fc:\n             lc = s.rfind(':', 0, lb)\n             if lc > fc:\n                ans = 4 + s[fc:lc].count('|')\nprint(ans)\n", "s = input()\n\ndef sovle(s):\n\n    i1 = s.find('[')\n    if i1 == -1:\n        return -1\n    s = s[i1+1:]\n    i2 = s.find(':')\n    if i2 == -1:\n        return -1\n\n    s = s[i2+1 :]\n    i1 = s.rfind(']')\n    if i1 == -1:\n        return -1\n    s = s[:i1]\n    i2 = s.rfind(':')\n    if i2 == -1:\n        return -1\n    s = s[:i2]\n    x = s.count('|')\n    return x+4\n\nprint(sovle(s))", "def solve(s):\n    if s.find('[') == -1:\n        return -1\n    s = s[s.find('['):]\n    #print(s)\n    if s.find(':') == -1:\n        return -1\n    s = s[s.find(':') + 1:]\n    #print(s)\n    if s.find(']') == -1:\n        return -1\n    s = s[:s.rfind(']')]\n    #print(s)\n    if s.find(':') == -1:\n        return -1\n    s = s[:s.rfind(':')]\n    #print(s)\n    return s.count('|') + 4\n\ns = input()\nprint(solve(s))", "s=input()\ni=s.find('[')\nif i==-1:\n    print(-1)\n    return\ns=s[i:]\ni=s.rfind(']')\n\nif i==-1:\n    print(-1)\n    return\ns=s[:i+1]\nl,h=0,0\nfor i,d in enumerate(s):\n    if d==':':\n        l=i\n        break\nfor i,d in enumerate(s):\n    if d==':':\n        h=i\nif l==h:\n    print(-1)\n    return\nc=0\nfor i in range(l+1,h):\n    if s[i]=='|':\n        c+=1\nprint(c+4)\n", "from sys import stdin\ns=stdin.readline().strip()\nx=-1\nfor i in range(len(s)):\n    if s[i]==\"[\":\n        x=i\n        break\ny=-1\nfor i in range(len(s)-1,-1,-1):\n    if s[i]==\"]\":\n        y=i\n        break\nif x==-1 or y==-1 or y<x:\n    print(-1)\n    return\nx1=-1\nfor i in range(x,y):\n    if s[i]==\":\":\n        x1=i\n        break\ny1=-1\nfor i in range(y-1,x,-1):\n    if s[i]==\":\":\n        y1=i\n        break\nif x1==-1 or y1==-1 or y1<=x1:\n    print(-1)\n    return\nans=4\nfor i in range(x1,y1):\n    if s[i]==\"|\":\n        ans+=1\nprint(ans)\n", "s = str(input().strip())\ni = 0\nn = len(s)\nwhile i < n and s[i] != '[':\n    i+=1\nif(i == n):\n    print(-1)\n    return\nj = n-1\nwhile j > i and s[j] != ']':\n    j-=1\nif(j <= i):\n    print(-1)\n    return\nwhile i < j and s[i] != ':':\n    i+=1\nif(i == j):\n    print(-1)\n    return\nwhile j > i and s[j] != ':':\n    j-=1\nif(j == i):\n    print(-1)\n    return\nk = i+1\nc = 0\nwhile k < j:\n    if(s[k] == '|'):\n        c+=1\n    k+=1\nprint(c+4)\n", "import sys\ns = input()\nl = len(s)\ns_list = [x for x in s]\n\ncounter = 0\ntry:\n\ta = s_list.index('[')\n\tcounter += a\n\ts_list = s_list[a + 1:]\nexcept:\n\tprint(-1)\n\treturn\n\ntry:\n\ta = s_list.index(':')\n\tcounter += a\n\ts_list = s_list[a + 1:]\nexcept:\n\tprint(-1)\n\treturn\n\ns_list_rev = s_list.copy()\ns_list_rev.reverse()\n\ntry:\n\tb = s_list_rev.index(']')\n\tcounter += b\n\ts_list_rev = s_list_rev[b+1:]\nexcept:\n\tprint(-1)\n\treturn\n\ntry:\n\tb = s_list_rev.index(':')\n\tcounter += b\n\ts_list_rev = s_list_rev[b+1:]\nexcept:\n\tprint(-1)\n\treturn\ns_list_rev = [x for x in s_list_rev if x != '|']\ncounter += len(s_list_rev)\nprint(l - counter)", "MOD = 10**9 + 7\nI = lambda:list(map(int,input().split()))\n\ns = input()\nres = 0\nn = len(s)\nst = -1\ne = -1\nfor i in range(n):\n    if s[i] == '[':\n        st = i\n        break\nfor i in range(n-1, -1, -1):\n    if s[i] == ']':\n        e = i\n        break\n# print(st , e)\nif st > e or st == -1 or e == -1:\n    print(-1)\n    return\na = -1\nb = -1\nfor i in range(st, e):\n    if s[i] == ':':\n        a = i\n        break\nfor i in range(e, st, -1):\n    if s[i] == ':':\n        b = i\n        break\nif a == b or a == -1 or b == -1:\n    print(-1)\n    return\ncount = 0\nfor i in range(a, b):\n    if s[i] == '|':\n        count += 1\nprint(4 + count)", "s=input()\nst=\"\"\nidx=-1\nfor i in range(len(s)):\n    if s[i]=='[':\n        idx=i\n        break\nif idx==-1:\n    print(-1)\n    return\nidxl=-1\nfor i in range(len(s)-1,-1,-1):\n    if s[i]==']' and i>idx:\n        idxl=i\n        break\nif idxl==-1:\n    print(-1)\n    return\ncol=col2=-1\nfor i in range(len(s)):\n    if s[i]==':' and i>idx and i<idxl:\n        col=i\n        break\nif col==-1:\n    print(-1)\n    return\nfor i in range(len(s)-1,-1,-1):\n    if s[i]==':' and i>col and i<idxl:\n        col2=i\n        break\nif col2==-1:\n    print(-1)\n    return\nans=0\nfor i in range(col+1,col2):\n    if s[i]=='|':\n        ans+=1\nprint(4+ans)\n    \n\n\n", "s = input()\nrev = s[::-1]\n\nleft = s.find(\"[\")\nif left != -1:\n    left = s.find(\":\", left)\n\nright = rev.find(\"]\")\nif right != -1:\n    right = rev.find(\":\", right)\n\nif left == -1 or right == -1:\n    print(-1)\n    return\nright = len(s)-right-1\nif left >= right:\n    print(-1)\n    return\n\nprint(4 + s[left:right].count(\"|\"))\n", "def ba(s):\n    c1 = s.find('[')\n    c2 = s.find(':', c1+1)\n    c3 = s.rfind(']', c2+1)\n    c4 = s.rfind(':', c2+1, c3)\n    if -1 in [c1, c2, c3, c4]:\n        return -1\n    return s.count('|', c2, c4)+4\n\n\nprint(ba(input()))\n\n", "s = input()\nif '[' in s and ']' in s:\n    a = s.index('[') + 1\n    b = len(s)-s[::-1].index(']') - 1\nelse:\n    print(-1)\n    return\ns = s[a:b]\nif s.count(':') >= 2:\n    a = s.index(':')+1\n    b = len(s)-s[::-1].index(':')-1\nelse:\n    print(-1)\n    return\nc = 0\nfor el in s[a:b]:\n    if el =='|':\n        c += 1\nprint(4 + c)", "s = input()\n\nb = [0]*len(s)\n\nob = 0\ncc = 0\np = -1\nq = -1\n\ncount = 0\n\nfor ind,c in enumerate(s):\n    if c == '[':\n        ob = 1\n    elif c == ':' and p >= 0:\n        q = ind\n    elif c == ':' and ob == 1 and p < 0:\n        p = ind\n    elif c == ']' and q >= 0:\n        cc = q\n    elif c == '|':\n        count += 1\n    b[ind] = count\n\nif cc > 0:\n    print( 4 + b[cc]-b[p])\nelse:\n    print(-1)\n", "s = input()\nif '[' in s and ']' in s and ':' in s:\n  e = s.count(':')\n  if e<2:\n    print(-1)\n  else:\n    a = s.index('[')\n    b = len(s)-1-s[::-1].index(']')\n    if b<a:\n      print(-1)\n    else:\n      if s[a+1:b].count(':')<2:\n        print(-1)\n      else:\n        st1 = True\n        count = 0\n        for i in range(a+1, b):\n          if st1 and s[i]==':':\n            pos1 = i\n            st1 = False\n          if s[i]==':':\n            pos2 = i\n      \n        for i in range(pos1+1, pos2):\n          if s[i]=='|':\n            count+=1\n      \n        print(count+4)\nelse:\n  print(-1) ", "s=input()\ni1=-1\ni2=-1\nk1=-1\nk2=-1\nc=0\nfor i in range(len(s)):\n  if(s[i]=='['):\n    i1=i\n    break\nfor i in range(len(s)-1,-1,-1):\n  if(s[i]==']'):\n    i2=i\n    break\nfor i in range(i1,i2+1):\n  if(s[i]==':'):\n    k1=i\n    break\nfor i in range(i2,i1-1,-1):\n  if(s[i]==':'):\n    k2=i\n    break\nfor i in range(k1,k2+1):\n  if(s[i]=='|'):\n    c+=1\n\nif(i1==-1 or i2==-1 or i1>=i2 or k1==-1 or k2==-1 or k1==k2):\n  print(-1)\nelse:\n  print(4+c)", "s = input()\nl = 0\nend = 0\ni = 1\n\nwhile i <= len(s):\n    if l == 0 and s[-i] == ']':\n        l += 1\n    elif l == 1 and s[-i] == ':':\n        l += 1\n        end = len(s) - i\n        break\n    i += 1\n\nif l < 2:\n    print(-1)\n    return\n\nfor i in range(0, end):\n    if l >= 4 and s[i] == '|':\n        l += 1\n    elif l == 2 and s[i] == '[':\n        l += 1\n    elif l == 3 and s[i] == ':':\n        l += 1\n\nif l >= 4:\n    print(l)\nelse:\n    print(-1)"]

test_examples/solutions_problem_2.json

@@ -0,0 +1 @@
+["num = list(map(int, input()))\nbest = num[:]\nfor i in range(-1, -len(num) - 1, -1):\n    if num[i] == 0:\n        continue\n    num[i] -= 1\n    for j in range(i + 1, 0):\n        num[j] = 9\n    if sum(num) > sum(best):\n        best = num[:]\ns = ''.join(map(str, best)).lstrip('0')\nprint(s)\n", "s_num = input()\nnum = int(s_num)\ndigs = [int(s_num[i]) for i in range(len(s_num))]\n\nmax_sum = sum(digs)\nres = num\nfor i in range(len(s_num)):\n    if (digs[i] != 0):\n        digs[i] -= 1\n        n_sum = sum(digs[:i + 1]) + 9 * (len(s_num) - i - 1)\n        if n_sum >= max_sum:\n            n_res = int(''.join([str(digs[i]) for i in range(i + 1)]) + '9' * (len(s_num) - i - 1))\n            if (n_sum == max_sum):\n                res = max(n_res, res)\n            else:\n                res = n_res\n            max_sum = n_sum\n\n        digs[i] += 1\nprint(res)\n", "a=int(input())\nif(a//10==0):\n    print(a)\n    return\nk=9\nwhile(k<a):\n    k=k*10+9\nif(k==a):\n    print(k)\nelse:\n    k//=10\n    k=int(str(a)[0]+str(k))\n    i=len(str(k))-1\n    z=k\n    while(z>a):\n        z=int(str(k)[0:i]+str(int(str(k)[i])-1)+str(k)[i+1:len(str(k))])\n        i-=1\n    print(z)    ", "x = int(input())\nif x < 10:\n    print(x)\nelif x == int(str(x)[0] + '9'*(len(str(x))-1)):\n    print(x)\nelse:\n    a = str(x)[0] + '9' * (len(str(x)) - 1)\n    a = list(a)\n    for i in range(len(a) - 1, -1, -1):\n        k = a[i]\n        a[i] = str(int(a[i]) - 1)\n        if x >= int(''.join(a)):\n            print(int(''.join(a)))\n            break\n        a[i] = k\n", "def sum_str(y):\n    return sum(map(int, str(y)))\n\n\nx = input()\nlength = len(x)\nbad_answer = str(int(x[0]) - 1) + '9' * (length - 1) \ntotal = sum_str(bad_answer)\n\n\nif length == 1 or sum_str(x) >= total:\n    print(x)\nelse:\n    for i in range(length - 1, 0, -1):\n        new_total = 9 * (length - i)\n        new_answer = str(int(x[:i]) - 1)\n        new_total += sum_str(new_answer)\n\n        if new_total >= total:\n            new_answer = new_answer if new_answer != '0' else ''\n            print(new_answer + '9' * (length - i))\n            break\n    else:\n        print(bad_answer)\n", "import sys\n\ndef calc(s):\n    res =0\n    for c in s:\n        res+= int(c)\n    return res\n\n\ns = list(sys.stdin.readline().rstrip())\nbest = \"\".join(s) \ncount = calc(s)\n\ni = len(s)-1\nwhile i!=0:\n    i-=1\n    if s[i+1]!= '9':\n        s[i+1] = '9'\n        while s[i]=='0':\n            s[i]='9'\n            i-=1\n        s[i] = chr(ord(s[i])-1)\n        c = calc(s)\n        if count < c:\n            count = c\n            best = \"\".join(s)\n\nif best[0] == '0':\n    best = best[1:]\n\nprint(best)", "x = input()\nn = len(x)\nif n == 1:\n    print(x)\n    return\nans = \"\"\ns = 0\nps = 0\npn = \"\"\nfor i in range(n):\n    ts = ps + int(x[i]) - 1 + 9 * (n - i - 1)\n    if ts >= s:\n        ans = pn + str(int(x[i]) - 1) + \"9\" * (n - i - 1)\n        s = ts\n    ps += int(x[i])\n    pn += x[i]\nif ps >= s:\n    ans = pn\nprint(int(ans))", "n = int(input())\n\ndef f(numb):\n    lst = [numb]\n    cap = 10\n\n    while numb // cap > 0:\n        lst.append((numb // cap - 1) * cap + cap - 1)\n        cap *= 10\n\n    return lst\n\ndef g(numb):\n    lst = []\n    while numb != 0:\n        lst.append(numb % 10)\n        numb //= 10\n\n    return lst\n\n\nmaximum = max([sum(g(i)) for i in f(n)])\n\nmaximum = [i for i in f(n) if maximum == sum(g(i))]\n\nprint(max(maximum))", "\"\"\" Created by Shahen Kosyan on 3/11/17 \"\"\"\n\ndef __starting_point():\n    x = input()\n\n    if int(x) < 10:\n        print(x)\n        return\n\n    arr = [int(a) for a in list(x)]\n    x_sum = sum(arr)\n\n    i = len(arr) - 1\n    answer = ''\n    while i > 0:\n        if arr[i] != 9 and arr[i] != 8:\n            arr[i - 1] -= 1\n            answer = '9' + answer\n        else:\n            change = False\n            for j in range(i - 1, 0, -1):\n                if arr[j] < 9:\n                    change = True\n                    break\n\n            if arr[i] == 8 and change:\n                answer = '9' + answer\n                arr[i - 1] -= 1\n            else:\n                if not change:\n                    answer = str(arr[i]) + answer\n                else:\n                    answer = '9' + answer\n\n        if i == 1 and arr[0] != 0:\n            answer = str(arr[0]) + answer\n        i -= 1\n\n    answer = [int(a) for a in list(answer)]\n    if x_sum == sum(answer):\n        print(x)\n    else:\n        answer = [str(a) for a in answer]\n        print(''.join(answer))\n\n__starting_point()", "x=input()\nl=len(x)\nx=int(x)\ns='9'*l\nsx=str(x)\nm=int(s)\nc=0\nwhile c!=1:\n    if m>x:\n        m=m-10**(l-1)\n    else:\n        c=1\nsm=str(m)\nmm=[] \nfor i in range(len(sm)):\n    mm.append(int(sm[i]))\nxx=[] \nfor i in range(l):\n    xx.append(int(sx[i]))\nif m==x:\n    print(m)\nelif sum(xx)==sum(mm):\n    print(x)\nelse:\n    k=len(xx)-1\n    while k>=0:\n        if sum(xx)<sum(mm):\n            if xx[k]==9:\n                k-=1\n            else:\n                xx[k]=9\n                xx[k-1]-=1\n                k-=1\n        else:\n            if xx[0]==0:\n                xx.remove(0)\n            for b in range(len(xx)):\n                xx[b]=str(xx[b])\n            ww=''.join(xx)\n            print(ww)\n            break", "x = input()\nvariants = [x] + [str(int(x[:i]) - 1) +\n    '9' * (len(x) - i) for i in range(1, len(x))]\nprint(int(max(variants, key=lambda x: (sum(map(int, x)), int(x)))))\n", "def sum_div(n):\n    summa = 0\n    while n > 0:\n        summa = summa + n % 10\n        n = n // 10\n    return summa\n\n\ndef run(n):\n    l_n = len(n)\n    left = ''\n    if l_n > 2 and '9' * l_n != n and n[1] == '9' and '9' * (l_n - 1) != n[1:]:\n        left = n[0]\n        n = n[1:]\n        while l_n > 1 and n[1] == '9':\n            left += n[1]\n            n = n[1:]\n            l_n = len(n)\n    l_n = len(n)\n    if len(n) == 1:\n        return n\n    elif '9' * (l_n - 1) == n[1:]:\n        return left + n\n    elif n[0] != '1':\n        min_number = int(str(int(n[0]) - 1) + '9' * (l_n - 1))\n        if sum_div(min_number) > sum_div(int(n)):\n            return left + str(min_number)\n        else:\n            return left + n\n    else:\n        min_number = int('9' * (l_n - 1)) if l_n > 1 else 0\n        if sum_div(min_number) > sum_div(int(n)):\n            return left + str(min_number)\n        else:\n            return left + n\n\n\nn = input()\nprint(run(n))\n", "#This code is dedicated to Olya S.\n\ndef e(x):\n    s=0\n    while x>0:\n        s+=x%10\n        x//=10\n    return s\n\ndef down(x):\n    l=len(x)-1\n    return str(int(x[0])-1)+'9'*l\n\nn=input()\nif len(n)>1 and n[1]=='9':\n    print(n[0],end='')\n    n=n[1:]\n    while len(n)>1 and n[0]=='9' and n[1]=='9':\n        print('9',end='')\n        n=n[1:]\n\nif e(int(n))>=e(int(down(n))):\n    print(n)\nelse:\n    print(int(down(n)))\n\n        \n     \n\n\n\n        \n\n", "def sum_n(n):\n    l = len(n)\n\n    summ = 0\n    for i in range(l):\n        summ += int(n[i])\n\n    return summ\n\ndef transfer(x, i):\n    x = list(x)\n    \n    x[i+1] = '9'\n    if x[i] != '0':\n        x[i] = str(int(x[i])-1)\n    else:\n        j = i\n        while (j > 0) and (int(x[j]) == 0):\n            x[j] = '9'\n            j -= 1\n        x[j] = str(int(x[j])-1)\n    if (x[0] == '0'):\n        del x[0]\n\n    return x\n\nx = list(input())\nmax_cifr = sum_n(x)\nmaxnum = x\nres = ''\n\nfor i in range(len(x)-2, -1, -1):\n    x = transfer(x, i)\n    if(max_cifr < sum_n(x)):\n        max_cifr = sum_n(x)\n        maxnum = x\n\nfor i in range(len(maxnum)):\n    res = res+maxnum[i]\n    \nprint(res)\n", "x = input()\nsum = 0\nfor i in x:\n    temp = int(i)\n    sum += temp\n\nxlen = len(x)\none = int(x[0])\ntry:\n    two = int(x[1])\nexcept:\n    two = 0\n\nif (two == 9):\n    count = 1\n    for i in range(1, xlen):\n        z = int(x[i])\n        if (z == 9):\n            count = i\n        else:\n            break\n    answ = x[0:count] + \"8\" + (\"9\" * (xlen - count - 1))\nelif (one == 1):\n    answ = '9' * (xlen - 1)\nelse:\n    answ = str((one - 1)) + (\"9\" * (xlen-1))\n\nansw = str(answ)\nsumansw = 0\nfor i in answ:\n    temp = int(i)\n    sumansw += temp\n\nif (sum >= sumansw):\n    print(x)\nelse:\n    print(answ)", "def sum1(x):  # \u043f\u043e\u0434\u0441\u0447\u0451\u0442 \u0441\u0443\u043c\u043c\u044b \u0446\u0438\u0444\u0440 \u0447\u0438\u0441\u043b\u0430 x\n    summa = 0\n    for i in x:\n        summa += int(i)\n    return summa\n\n\nx = input()\nc = sum1(x)\nresult = int(x)\nn = len(x) - 1\nj = n\nfor i in range(0, n):\n    if x[i] != '0':\n        ni = int(x[i]) - 1  # \u0443\u043c\u0435\u043d\u044c\u0448\u0430\u044e i-\u044b\u0439 \u0440\u0430\u0437\u0440\u044f\u0434 \u043d\u0430 1\n        xi = x[0:i] + str(ni) + '9' * j  # \u0441\u0442\u0440\u043e\u044e \u043d\u043e\u0432\u043e\u0435 \u0447\u0438\u0441\u043b\u043e\n        j -= 1\n        ci = sum1(xi)\n        if c < ci:\n            c = ci\n            result = int(xi)\n        elif c == ci and result < int(xi):\n            result = int(xi)\n    else:\n        j -= 1\n        continue\nprint(result)\n", "def f(n, k):\n    n = str(n)\n    if n[k] == \"0\":\n        return f(n, k - 1)\n    a = []\n    for i in n:\n        a.append(int(i))\n    n = a\n    n[k] = int(n[k]) - 1\n    n[k + 1::] = [9] * (len(n) - k - 1)\n    return n\na = input()\nn = len(a)\nans = [int(x) for x in a]\nms = sum(ans)\nfor i in range(0, n):\n    ca = f(a, i)\n    cs = sum(ca)\n    if cs> ms:\n        ans = ca\n        ms = cs\n    elif cs == ms:\n        if int(''.join([str(_) for _ in ca])) > int(''.join([str(_) for _ in ans])):\n            ans = ca\nprint(int(''.join([str(_) for _ in ans])))", "n = int(input().strip())\n\ns = []\nwhile n > 0:\n    s.append(n % 10)\n    n //= 10\ns = s[::-1]\n\nn = len(s)\nans = 0\nbest = -1\nfor i in range(n):\n    res = sum(s[:i + 1]) - 1 + 9 * (n - i - 1)\n    if res >= ans:\n        ans = res\n        best = i\n\ndef get(s, pos):\n    ans = 0\n    for i in range(len(s)):\n        if i > pos:\n            ans = ans * 10 + 9\n        else:\n            ans = ans * 10 + s[i]\n            if i == pos:\n                ans -= 1\n    return ans\n\nif sum(s) >= ans:\n    print(get(s, n))\nelse:\n    print(get(s, best))\n\n", "def main():\n\n\tdef sum(x):\n\t\tres = 0\n\n\t\twhile x > 0:\n\t\t\tres += x % 10\n\t\t\tx //= 10\n\n\t\treturn res\n\n\tn = input()\n\tfirst = n[0]\n\tp = [1]\n\n\tfor i in range(1, 20):\n\t\tp.append(p[-1] * 10)\n\n\tdata = []\t\n\tfor i in range(len(n)):\n\t\tif i > 0 and n[i] == '0':\n\t\t\tcontinue\n\t\ttemp = n[:i] + str(max(0, int(n[i]) - 1)) + \"9\"* (len(n) - i - 1)\n\t\tdata.append((sum(int(temp)), int(temp)))\n\n\tdata.append((sum(int(n)), int(n)))\n\t\n\tdata.sort(reverse=True)\n\n\tprint(data[0][1])\n\n\treturn\n\ndef __starting_point():\n\tmain()\n__starting_point()", "def cnt_sum(str_num):\n\tsum = 0\n\tfor a in str_num:\n\t\tsum += ord(a) - ord('0')\n\treturn sum\n\nstr_a = input().strip()\nmax_sum = cnt_sum(str_a)\nans = str_a\ncnt_digit = len(str_a)\n\nfor i in range(cnt_digit - 1, -1, -1):\n\tif str_a[i] != '0':\n\t\tnew_str = str_a[:i] + chr(ord(str_a[i]) - 1) + '9'*(cnt_digit - i - 1)\n\t\tcur_sum = cnt_sum(new_str)\n\t\tif cur_sum > max_sum:\n\t\t\tmax_sum = cur_sum\n\t\t\tans = new_str\n\nprint(int(ans))\n", "def summaX(x):\n    k=0\n    for el in x:\n        k+=int(el)\n    return k\nn=input();N=[];Z=[]\nfor el in n:\n    N.append(el)\nz=summaX(N)\nZ=N.copy()\nfor i in range(1,len(N)):\n    if int(N[i])!=9:\n        N[i-1]=int(N[i-1])-1\n        for j in range(i,len(n)):\n            N[j]=9\nif z>=summaX(N):\n    for el in Z:\n        print(el,end='')\nelse:\n    if N[0]==0:\n        N.pop(0)\n    for el in N:\n        print(el,end='')\n", "n = int(input())\n\ndef sumd(n):\n\tj = n\n\tsumn = 0\n\twhile j:\n\t\tsumn += j % 10\n\t\tj //= 10\n\treturn sumn\n\nj = n\nstrn = str(n)\nl = len(strn)\nsumn = sumd(n)\n\nstra = [i for i in str(n)]\ni = 1\nwhile i < l and stra[i] == '9':\n\ti += 1\nif (i != l):\n\tstra[i - 1] = str(int(stra[i - 1]) - 1)\n\twhile i < l:\n\t\tstra[i] = '9'\n\t\ti += 1\n\nss = ''\nfor i in range(l):\n\tss += stra[i]\nif ss[0] == '0':\n\tss = ss[1:]\nsn = int(ss)\n\nif sn < n and sumd(sn) <= sumn:\n\tss = strn\n\tsn = n\n\nprint(ss)\n", "from random import randint\n\ndef f(s):\n    a = 0\n    for i in s:\n        a += int(i)\n    return a\n\ndef solve(n):\n    n1 = list(str(n))\n    ans = 0\n    maxx = 0\n    for i in range(len(n1)):\n        n2 = n1[:i] + [str(int(n1[i]) - 1)] + ['9' for j in range(len(n1) - i - 1)]\n        if f(n2) >= maxx:\n            maxx = f(n2)\n            ans = n2\n    if f(n1) >= maxx:\n        maxx = f(n1)\n        ans = n1\n    return [int(''.join(ans)), maxx]\n\ndef tl(n):\n    ans = 0\n    maxx = 0\n    for i in range(1, n + 1):\n        if f(list(str(i))) >= maxx:\n            maxx = f(list(str(i)))\n            ans = i\n    return [ans, maxx]\n\n'''for kkk in range(100):\n    n = randint(1, 10 ** 5)\n    c1 = solve(n)\n    c2 = tl(n)\n    if c1 != c2:\n        print(n)\n        print(c1)\n        print(c2)\nprint('ok')'''\nn = int(input())\nprint(solve(n)[0])\n", "a = [1, 2, 3, 4, 5, 6, 7, 8, 9]\nfor length in range(2, 30):\n    for first in range(1, 10):\n        for pos in range(1, length):\n            a.append(int(str(first) + '9' * (pos - 1) + '8' + '9' * (length - pos - 1)))\n        a.append(int(str(first) + '9' * (length - 1)))\n    \nn = int(input())\nl = 0\nr = len(a)\nwhile l < r - 1:\n    middle = (l + r) // 2\n    if (a[middle] <= n):\n        l = middle\n    else:\n        r = middle\n        \nprint(a[l])", "def get(s):\n    ans = 0\n    for i in s:\n        ans += (ord(i) - ord('0'))\n    return ans\n\n\ndef solve1():\n    x = input()\n    n = len(x)\n    best_ans = x\n    best_val = get(x)\n    ans = str('' if int(x[0]) - 1 == 0 else int(x[0]) - 1) + '9' * (n - 1)\n    if get(ans) > best_val or (get(ans) >= best_val and int(ans) > int(best_ans)):\n        best_ans = ans\n        best_val = get(ans)\n    for i in range(1, n):\n        #print(ans)\n        ans = x[:i] + str(int(x[i]) - 1) + '9' * (n - i - 1)\n        if get(ans) > best_val or (get(ans) >= best_val and int(ans) > int(best_ans)):\n            best_ans = ans\n            best_val = get(ans)\n    return best_ans\n    \nbest = [0] * 10000\ndef solve2():\n    nonlocal best\n    was = 0\n    for i in range(1, 10000):\n        if get(str(i)) >= was:\n            best[i] = i\n            was = get(str(i))\n        else:\n            best[i] = best[i - 1]\n    \ndef stress():\n    solve2()\n    for i in range(1, 10000):\n        if int(solve1(str(i))) != best[i]:\n            print(i, best[i], solve1(str(i)))\n\n#stress()\nprint(solve1())"]

testing_util.py

@@ -0,0 +1,525 @@
+import json
+import sys
+import faulthandler
+import platform
+
+# used for debugging to time steps
+from datetime import datetime
+
+# to run the solution files we're using a timing based approach
+import signal
+
+import numpy as np
+# for capturing the stdout
+from io import StringIO
+# used for testing the code that reads from input
+from unittest.mock import patch, mock_open
+
+from pyext import RuntimeModule
+
+from enum import Enum
+class CODE_TYPE(Enum):
+    call_based = 0
+    standard_input = 1
+
+# stuff for setting up signal timer
+class TimeoutException(Exception):
+    pass
+def timeout_handler(signum, frame):
+    print("alarm went off")
+    #return
+    raise TimeoutException
+signal.signal(signal.SIGALRM, timeout_handler)
+timeout = 4  # seconds
+
+# used to capture stdout as a list
+# from https://stackoverflow.com/a/16571630/6416660
+# alternative use redirect_stdout() from contextlib
+class Capturing(list):
+    def __enter__(self):
+        self._stdout = sys.stdout
+        sys.stdout = self._stringio = StringIO()
+        # Make closing the StringIO a no-op
+        self._stringio.close = lambda x: 1
+        return self
+    def __exit__(self, *args):
+        self.extend(self._stringio.getvalue().splitlines())
+        del self._stringio    # free up some memory
+        sys.stdout = self._stdout
+
+
+def run_test(sample, test=None, debug=False):
+    """
+    if test(generated_code) is not None it'll try to run the code.
+    otherwise it'll just return an input and output pair.
+    """
+    # Disable functionalities that can make destructive changes to the test.
+    reliability_guard()
+    
+    if debug:
+        print(f"start = {datetime.now().time()}")
+
+    try:
+        in_outs = json.loads(sample["input_output"])
+    except ValueError:
+        in_outs = None
+    if in_outs:
+        if in_outs.get("fn_name") is None:
+            which_type = CODE_TYPE.standard_input  # Standard input
+            method_name = None
+        else:
+            which_type = CODE_TYPE.call_based  # Call-based
+            method_name = in_outs["fn_name"]
+
+    if debug:
+        print(f"loaded input_output = {datetime.now().time()}")
+ 
+    if test is None:
+        return in_outs
+    elif test is not None:
+        results = []
+        sol = "import sys\nimport time\nimport itertools\nfrom itertools import accumulate, product, permutations, combinations\nimport collections\nfrom collections import Counter, OrderedDict, deque, defaultdict, ChainMap\nfrom functools import lru_cache\nimport math\nfrom math import sqrt, sin, cos, tan, ceil, fabs, floor, gcd, exp, log, log2\nimport fractions\nfrom typing import List, Tuple\nimport numpy as np\nimport random\nimport heapq\nfrom heapq import *\n"
+        if debug:
+            print(f"loading test code = {datetime.now().time()}")
+ 
+        if which_type == CODE_TYPE.call_based:
+            sol += test
+            if debug:
+                print(f"sol = {sol}")
+            signal.alarm(timeout)
+            try:
+                tmp_sol = RuntimeModule.from_string("tmp_sol", "", sol)
+                if "class Solution" not in test:
+                    tmp = tmp_sol
+                else:
+                    tmp = tmp_sol.Solution()
+                signal.alarm(0)
+            except Exception as e:
+                signal.alarm(0)
+                if debug:
+                     print(f"type 0 compilation error = {e}")
+                results.append(-2)
+                return results
+            signal.alarm(0)
+
+        elif which_type == CODE_TYPE.standard_input:
+            # sol
+            tmp_test = test.split("\n")
+
+            new_test = []
+            for x in tmp_test:
+                if (not x.startswith("from ")) and (not x.startswith("import ")):
+                    new_test.append("\t" + x + "\n")
+                else:
+                    new_test.append(x + "\n")
+            tmp_test = new_test
+            
+            new_test = ""
+            started = False
+            for i in tmp_test:
+                if i.startswith("\t") and not started:
+                    new_test += "stdin = sys.stdin\nstdout = sys.stdout\n"
+                    new_test += "def code():\n"
+                    new_test += i
+                    started = True
+                elif started and ((i.startswith("from ")) or (i.startswith("import "))): 
+                    new_test += "\t" + i
+                else:
+                    new_test += i
+            tmp_test = new_test
+
+            sol += tmp_test
+            if debug:
+                print(f"sol = {sol}")
+            method_name = "code"
+            signal.alarm(timeout)
+            try:
+                tmp_sol = RuntimeModule.from_string("tmp_sol", "", sol)
+                tmp = tmp_sol
+                signal.alarm(0)
+            except Exception as e:
+                signal.alarm(0)
+                if debug:
+                    print(f"type 1 compilation error = {e}")
+                results.append(-2)
+                return results
+            signal.alarm(0)
+        if debug:
+            print(f"get method = {datetime.now().time()}")
+ 
+        try:
+            method = getattr(tmp, method_name)  # get_attr second arg must be str
+        except:
+            signal.alarm(0)
+            e = sys.exc_info()
+            print(f"unable to get function error = {e}")
+            results.append(-2)
+            return results
+
+        for index, inputs in enumerate(in_outs["inputs"]):
+            # JSON forces dictionaries to have string keys; this undoes this (assuming a singleton list)
+            try:
+                if isinstance(inputs[0], dict):
+                    inputs = [{int(k): v for k,v in inputs[0].items()}]
+            except:
+                True
+            try:
+                if isinstance(in_outs["outputs"][index], dict):
+                    in_outs["outputs"][index] = [{int(k): v for k,v in in_outs["outputs"][index].items()}]
+            except:
+                True
+            try:
+                if isinstance(in_outs["outputs"][index][0], dict):
+                    in_outs["outputs"][index] = [{int(k): v for k,v in in_outs["outputs"][index][0].items()}]
+            except:
+                True
+
+            if debug:
+                print(f"time: {datetime.now().time()} testing index = {index}  inputs = {inputs}, {type(inputs)}. type = {which_type}")
+            if which_type == CODE_TYPE.call_based:  # Call-based
+                signal.alarm(timeout)
+                faulthandler.enable()
+                try:
+                    output = method(*inputs)
+
+                    # ground truth sequences are not tuples
+                    if isinstance(output, tuple):
+                        output = list(output)
+                    
+                    tmp_result = output == in_outs["outputs"][index]
+                    if isinstance(in_outs["outputs"][index], list) and in_outs["outputs"][index]:
+                        tmp_result = tmp_result or (output == in_outs["outputs"][index][0])
+
+                    # ground truth sequences are not tuples
+                    try:
+                        if isinstance(output[0], tuple):
+                            tmp_result = tmp_result or ([list(x) for x in output] == in_outs["outputs"][index][0])
+                    except:
+                        True
+                    results.append(tmp_result)
+
+                    # reset the alarm
+                    signal.alarm(0)
+                except Exception as e:
+                    signal.alarm(0)
+                    faulthandler.disable()
+                    if debug:
+                        print(f"Standard input runtime error or time limit exceeded error = {e}")
+                    results.append(-1)
+                    continue
+                faulthandler.disable()
+                signal.alarm(0)
+                if debug:
+                    print(f"outputs = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}")
+            elif which_type == CODE_TYPE.standard_input:  # Standard input
+                faulthandler.enable()
+                signal.alarm(timeout)
+                passed = False
+
+                if isinstance(inputs, list):
+                    inputs = "\n".join(inputs)
+                if isinstance(in_outs['outputs'][index], list):
+                    in_outs['outputs'][index] = "\n".join(in_outs['outputs'][index])
+
+                with Capturing() as output:
+                    try:
+                        call_method(method, inputs)
+                        # reset the alarm
+                        signal.alarm(0)
+                        passed = True
+                    except Exception as e:
+                        # runtime error or took too long
+                        signal.alarm(0)
+                        print(f"Call-based runtime error or time limit exceeded error = {repr(e)}{e}")
+                        results.append(-1)
+                    signal.alarm(0)
+
+                if not passed:
+                    if debug:
+                        nl = "\n"
+                        if not isinstance(inputs, list):
+                            print(f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}")
+                        else:
+                            print(f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}")
+                    continue
+
+                if passed and debug:
+                    print(f"==> output = {output}, test outputs = {in_outs['outputs'][index]}")
+
+                if custom_compare_(output, in_outs['outputs'][index]):
+                    tmp_result = True
+                    results.append(tmp_result)
+                    continue
+
+                # ground truth sequences are expressed as lists not tuples
+                if isinstance(output, tuple):
+                    output = list(output)
+
+                tmp_result = False
+                try:
+                    tmp_result = (output == [in_outs["outputs"][index]])
+                    if isinstance(in_outs["outputs"][index], list):
+                        tmp_result = tmp_result or (output == in_outs["outputs"][index])
+                        if isinstance(output[0], str):
+                            tmp_result = tmp_result or ([e.strip() for e in output] == in_outs["outputs"][index])
+                except Exception as e:
+                    if debug:
+                        print(f"Failed check1 exception = {e}")
+                    pass
+
+                if tmp_result == True:  
+                    results.append(tmp_result)
+                    continue
+
+                # try one more time without \n
+                if isinstance(in_outs["outputs"][index], list):
+                    for tmp_index, i in enumerate(in_outs["outputs"][index]):
+                        in_outs["outputs"][index][tmp_index] = i.split("\n")
+                        in_outs["outputs"][index][tmp_index] = [x.strip() for x in in_outs["outputs"][index][tmp_index] if x]
+                else:
+                    in_outs["outputs"][index] = in_outs["outputs"][index].split("\n")
+                    in_outs["outputs"][index] = list(filter(len, in_outs["outputs"][index]))
+                    in_outs["outputs"][index] = list(map(lambda x:x.strip(), in_outs["outputs"][index]))
+
+                try:
+                    tmp_result = (output == [in_outs["outputs"][index]])
+                    if isinstance(in_outs["outputs"][index], list):
+                        tmp_result = tmp_result or (output == in_outs["outputs"][index])
+                except Exception as e:
+                    if debug:
+                        print(f"Failed check2 exception = {e}")
+                    pass
+
+                if tmp_result == True:
+                    results.append(tmp_result)
+                    continue
+
+                # try by converting the output into a split up list too
+                if isinstance(output, list):
+                    output = list(filter(len, output))
+
+                if debug:
+                    nl = "\n"
+                    if not isinstance(inputs, list):
+                        print(f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}") 
+                    else:
+                        print(f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}") 
+                
+                if tmp_result == True:
+                    results.append(tmp_result)
+                    continue
+
+                try:
+                    tmp_result = (output == [in_outs["outputs"][index]])
+                    if isinstance(in_outs["outputs"][index], list):
+                        tmp_result = tmp_result or (output == in_outs["outputs"][index])
+                except Exception as e:
+                    if debug:
+                        print(f"Failed check3 exception = {e}")
+                    pass
+
+                try:
+                    output_float = [float(e) for e in output]
+                    gt_float = [float(e) for e in in_outs['outputs'][index]]
+                    tmp_result = tmp_result or ((len(output_float) == len(gt_float)) and np.allclose(output_float, gt_float))
+                except Exception as e:
+                    pass
+                try:
+                    if isinstance(output[0], list):
+                        output_float = [float(e) for e in output[0]]
+                        gt_float = [float(e) for e in in_outs['outputs'][index][0]]
+                        tmp_result = tmp_result or ((len(output_float) == len(gt_float)) and np.allclose(output_float, gt_float))
+                except Exception as e:
+                    pass
+
+                if tmp_result == True:
+                    results.append(tmp_result)
+                    continue
+
+                # try by converting the stuff into split up list
+                if isinstance(in_outs["outputs"][index], list):
+                    for tmp_index, i in enumerate(in_outs["outputs"][index]):
+                        in_outs["outputs"][index][tmp_index] = set(i.split())
+                else:
+                    in_outs["outputs"][index] = set(in_outs["outputs"][index].split())
+
+                try:
+                    tmp_result = (output == in_outs["outputs"][index])
+                except Exception as e:
+                    if debug:
+                        print(f"Failed check4 exception = {e}")
+                    continue
+
+                if tmp_result == True:
+                    results.append(tmp_result)
+                    continue 
+
+                # try by converting the output into a split up list too
+                if isinstance(output, list):
+                    for tmp_index, i in enumerate(output):
+                        output[tmp_index] = i.split()
+                    output = list(filter(len, output))
+                    for tmp_index, i in enumerate(output):
+                        output[tmp_index] = set(i)    
+                else:
+                    output = output.split()
+                    output = list(filter(len, output))
+                    output = set(output)
+
+                try:
+                    tmp_result = (set(frozenset(s) for s in output) == set(frozenset(s) for s in in_outs["outputs"][index]))
+                except Exception as e:
+                    if debug:
+                        print(f"Failed check5 exception = {e}")
+
+
+                # if they are all numbers, round so that similar numbers are treated as identical
+                try:
+                    tmp_result = tmp_result or (set(frozenset(round(float(t),3) for t in s) for s in output) ==\
+                        set(frozenset(round(float(t),3) for t in s) for s in in_outs["outputs"][index]))
+                except Exception as e:
+                    if debug:
+                        print(f"Failed check6 exception = {e}")
+                
+                if tmp_result == True and debug:
+                    print("PASSED")
+ 
+                results.append(tmp_result)
+                
+                if debug:
+                    nl = "\n"
+                    if not isinstance(inputs, list):
+                        print(f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}")
+                    else:
+                        print(f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}") 
+
+
+    return results
+
+
+def custom_compare_(output, ground_truth):
+    
+    if isinstance(output, list):
+        output_1 = "\n".join(output)
+        if stripped_string_compare(output_1, ground_truth):
+            return True
+
+    if isinstance(output, list):
+        output_2 = [o.lstrip().rstrip() for o in output]
+        output_2 = "\n".join(output_2)
+        if stripped_string_compare(output_2, ground_truth):
+            return True
+
+    return False
+
+def stripped_string_compare(s1, s2):
+    s1 = s1.lstrip().rstrip()
+    s2 = s2.lstrip().rstrip()
+    return s1 == s2
+
+def call_method(method, inputs):
+
+    if isinstance(inputs, list):
+        inputs = "\n".join(inputs)
+
+    inputs_line_iterator = iter(inputs.split("\n"))
+
+    # sys.setrecursionlimit(10000)
+
+    # @patch('builtins.input', side_effect=inputs.split("\n"))
+    @patch('builtins.open', mock_open(read_data=inputs))
+    @patch('sys.stdin', StringIO(inputs))
+    @patch('sys.stdin.readline', lambda *args: next(inputs_line_iterator))
+    @patch('sys.stdin.readlines', lambda *args: inputs.split("\n"))
+    @patch('sys.stdin.read', lambda *args: inputs)
+    # @patch('sys.stdout.write', print)
+    def _inner_call_method(_method):
+        try:
+            return _method()
+        except SystemExit as e:
+            pass
+        finally:
+            pass
+    return _inner_call_method(method) 
+
+
+
+
+def reliability_guard(maximum_memory_bytes=None):
+    """
+    This disables various destructive functions and prevents the generated code
+    from interfering with the test (e.g. fork bomb, killing other processes,
+    removing filesystem files, etc.)
+    WARNING
+    This function is NOT a security sandbox. Untrusted code, including, model-
+    generated code, should not be blindly executed outside of one. See the
+    Codex paper for more information about OpenAI's code sandbox, and proceed
+    with caution.
+    """
+
+    if maximum_memory_bytes is not None:
+        import resource
+
+        resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes))
+        resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes))
+        if not platform.uname().system == "Darwin":
+            resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes))
+
+    faulthandler.disable()
+
+    import builtins
+
+    builtins.exit = None
+    builtins.quit = None
+
+    import os
+
+    os.environ["OMP_NUM_THREADS"] = "1"
+
+    os.kill = None
+    os.system = None
+    os.putenv = None
+    os.remove = None
+    os.removedirs = None
+    os.rmdir = None
+    os.fchdir = None
+    os.setuid = None
+    os.fork = None
+    os.forkpty = None
+    os.killpg = None
+    os.rename = None
+    os.renames = None
+    os.truncate = None
+    os.replace = None
+    os.unlink = None
+    os.fchmod = None
+    os.fchown = None
+    os.chmod = None
+    os.chown = None
+    os.chroot = None
+    os.fchdir = None
+    os.lchflags = None
+    os.lchmod = None
+    os.lchown = None
+    os.getcwd = None
+    os.chdir = None
+
+    import shutil
+
+    shutil.rmtree = None
+    shutil.move = None
+    shutil.chown = None
+
+    import subprocess
+
+    subprocess.Popen = None  # type: ignore
+
+    __builtins__["help"] = None
+
+    import sys
+
+    sys.modules["ipdb"] = None
+    sys.modules["joblib"] = None
+    sys.modules["resource"] = None
+    sys.modules["psutil"] = None
+    sys.modules["tkinter"] = None

tests.py

@@ -0,0 +1,14 @@
+import json
+from evaluate import load
+
+solution_sample1 = json.load(open("test_examples/solutions_problem_1.json", "r"))
+solution_sample2 = json.load(open("test_examples/solutions_problem_2.json", "r"))
+single_solutions = [solution_sample1[:1], solution_sample2[:1]]
+multiple_solutions = [solution_sample1[:3], solution_sample2[:3]]
+
+metric = load("codeparrot/apps_metric")
+result_1 = metric.compute(predictions=single_solutions, level="all")
+result_2 = metric.compute(predictions=multiple_solutions, level="all", k_list=[1, 2, 3])
+
+assert result_1 == {'avg_accuracy': 1.0, 'strict_accuracy': 1.0, 'pass_at_k': None}
+assert result_2 == {'avg_accuracy': None, 'strict_accuracy': None, 'pass_at_k': {'pass@1': 1.0, 'pass@2': 1.0, 'pass@3': 1.0}}

utils.py

@@ -0,0 +1,212 @@
+import itertools
+import json
+import multiprocessing
+import numpy as np
+from typing import Dict
+from datasets import load_dataset
+from .testing_util import run_test
+
+DATASET = "codeparrot/apps"
+TIMEOUT = 10
+
+def check_correctness(sample, generation, timeout, debug=True):
+    """Check correctness of code generation with a global timeout.
+    The global timeout is to catch some extreme/rare cases not handled by the timeouts
+    inside `run_test`"""
+    def _temp_run(sample, generation, debug, result):
+        result.append(run_test(sample, test=generation, debug=debug))
+
+    manager = multiprocessing.Manager()
+    result = manager.list()
+    p = multiprocessing.Process(target=_temp_run, args=(sample, generation, debug, result))
+    p.start()
+    p.join(timeout=timeout + 1)
+    if p.is_alive():
+        p.kill()
+    if not result:
+        in_outs = json.loads(sample["input_output"])
+        # consider that all tests failed
+        result = [[-1 for i in range(len(in_outs["inputs"]))]]
+        if debug:
+            print(f"global timeout")
+    return result[0]
+
+
+def evaluate_generations(generations: list, level: str = "all", debug: bool = False):
+    """We take the list of code generations and try to compile them
+     and the run their corresponding unit tests which are retrieved from the APPS dataset.
+
+    Args:
+        generations: list of code generations (same order as samples in APPS dataset)
+        level: difficulty level used in the generation, can be "all", "introductory", "interview" or "competition"
+
+    Returns:
+        results: dictionary of results, key is the problem index, value is a list of results for each generation
+        [-2] = compile error, [-1] = runtime error [False] = failed test case [True] = passed test case
+     """
+
+    # generations are code generations in the same order of the dataset
+    apps_eval = load_dataset(DATASET, split="test", difficulties=[level])
+    results = {}
+    for index in range(len(generations)):
+        # code generations for problem (index)
+        problem_generations = generations[index]
+        # get corresponding samples from APPS dataset
+        sample = apps_eval[index]
+        res = []
+        # loop over the generations
+        for o_idx, o in enumerate(problem_generations):
+            curr_res = [-2]
+            try:
+                curr_res = check_correctness(sample, o, timeout=TIMEOUT, debug=debug)
+                if debug:
+                    print(f"\nSuccessful compilation of task {index}!")
+                fixed = []
+                for e in curr_res:
+                    if isinstance(e, np.ndarray):
+                       e = e.item(0)
+                    if isinstance(e, np.bool_):
+                        e = bool(e)
+                    fixed.append(e)
+                curr_res = fixed
+                if not np.all(curr_res):
+                    if debug:
+                        print(f"Results were not True for all test cases")
+            except Exception as e:
+                if debug:
+                    print(f"Compilation failed, test framework exception = {repr(e)}{e}\n")
+                break
+            finally:
+                assert isinstance(curr_res, list)
+                res.append(curr_res)
+        results[index] = res
+    return results
+
+
+def estimate_pass_at_k(num_samples, num_correct, k):
+    """Estimates pass@k of each problem and returns them in an array."""
+
+    def estimator(n: int, c: int, k: int) -> float:
+        """Calculates 1 - comb(n - c, k) / comb(n, k)."""
+        if n - c < k:
+            return 1.0
+        return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1))
+
+    if isinstance(num_samples, int):
+        num_samples_it = itertools.repeat(num_samples, len(num_correct))
+    else:
+        assert len(num_samples) == len(num_correct)
+        num_samples_it = iter(num_samples)
+
+    return np.array([estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)])
+
+
+def get_results(results: Dict[int, list], count_errors: bool = False, k_list: list = [1, 10, 100]):
+    """
+    Given the results evaluated against the testcases we output some statistics.
+    For single generations:
+    >>> example_results = {0: [[-2]], 1: [[False,False]], 2: [[True,True]], 3: [[False,True,False,True]], 4: [[-1,-1]]}
+    >>> get_results(example_results, count_errors=True)
+    Computing accuracy metrics...
+    number of compile errors = 1 avg = 0.2
+    number of runtime errors = 1 avg = 0.2
+    number of problems evaluated = 5
+    Average Accuracy : 0.3
+    Strict Accuracy : 0.2
+    {'avg_accuracy': 0.3, 'strict_accuracy': 0.2, 'pass_at_k': None}
+
+    For multiple generations:
+    >>> example_results = {0: [[-2], [True, True, True]], 1: [[-1,-1, -1], [True, False, True]]}
+    >>> get_results(example_results, k_list=[1, 2])
+    Computing pass@k metric for multiple generations...
+    {'pass@1': 0.25, 'pass@2': 0.5}
+    {'avg_accuracy': None, 'strict_accuracy': None, 'pass_at_k': {'pass@1': 0.25, 'pass@2': 0.5}}
+    """
+
+    metrics = {"avg_accuracy": None, "strict_accuracy": None, "pass_at_k": None}
+
+    if len(results[0]) == 1:
+        # for single generations we compute average accuracy and stric accuracy: original APPS metrics
+        print("Computing accuracy metrics...")
+        res = []
+        per_prob_res = []
+        all_correct = []
+        for index in results:
+            problem_results = np.asarray(results[index])
+            res.extend(problem_results)
+            per_prob_res.append(np.mean(problem_results > 0))
+            all_correct.append(np.all(problem_results > 0))
+        # we count campilation and runtime errors once per pronlem
+        compile_errors = len([e for e in res if -2 in e])
+        runtime_errors = len([e for e in res if -1 in e])
+        total_testcases = len(res)
+        if count_errors:
+            print(f"number of compile errors = {compile_errors} avg = {compile_errors / total_testcases}")
+            print(f"number of runtime errors = {runtime_errors} avg = {runtime_errors / total_testcases}")
+            print(f"number of problems evaluated = {total_testcases}")
+
+        print(f"Average Accuracy : {np.mean(per_prob_res)}")
+        print(f"Strict Accuracy : {np.mean(all_correct)}")
+        metrics["avg_accuracy"] = np.mean(per_prob_res)
+        metrics["strict_accuracy"] = np.mean(all_correct)
+
+    else:
+        # for multiple generations we use pass@k metric used in the HumanEval benchmark
+        # we use strict accuracy, a generation is valid if it has to pass all the tests
+        print("Computing pass@k metric for multiple generations...")
+        # total is list with nb generations per task (task=index)
+        # correct is number of generations that passed all tests per task
+        total = []
+        correct = [] 
+        for index in results:
+            all_correct = []
+            for generation in results[index]:
+                gen = np.array(generation)
+                all_correct.append(np.all(gen>0))
+            total.append(len(all_correct))
+            correct.append(sum(all_correct))
+        total = np.array(total)
+        correct = np.array(correct)
+        ks = k_list
+        pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean() for k in ks if (total >= k).all()}
+        print(pass_at_k)
+        metrics["pass_at_k"] = pass_at_k
+    return metrics
+
+def compute_metrics(generations, level="all", k_list=[1, 10, 100], count_errors=True, debug=False):
+    """Return metrics for the given generations.
+    Args:
+        generations: list of code generations for each problem (each generation is a list of generations)
+        k_list: list of k values to compute pass@k when using multiple generations
+        count_errors: whether to count compilation and runtime errors when using single generations
+        level: difficulty level in APPS dataset that was used for the given generations (from: "all", "introductory", "interview", "competition")
+    Returns:
+        metrics: dict of metrics  
+
+    Examples:
+
+    >>> import json
+    >>> # lists of solutions to the two first APPS problems (note not all solutions pass all tests)
+    >>> solution_sample1 = json.load(open("test_examples/solutions_problem_1.json", "r"))
+    >>> solution_sample2 = json.load(open("test_examples/solutions_problem_2.json", "r"))
+    >>> single_solutions = [solution_sample1[:1], solution_sample2[:1]]
+    >>> compute_metrics(single_solutions, level="all")
+    Computing accuracy metrics...
+    number of compile errors = 0 avg = 0.0
+    number of runtime errors = 0 avg = 0.0
+    number of problems evaluated = 2
+    Average Accuracy : 1.0
+    Strict Accuracy : 1.0
+    {'avg_accuracy': 1.0, 'strict_accuracy': 1.0, 'pass_at_k': None}
+    >>> multiple_solutions = [solution_sample1[:3], solution_sample2[:3]]
+    >>> compute_metrics(multiple_solutions, level="all", k_list=[1, 2, 3])
+    Computing pass@k metric for multiple generations...
+    {'pass@1': 1.0, 'pass@2': 1.0, 'pass@3': 1.0}
+    {'avg_accuracy': None, 'strict_accuracy': None, 'pass_at_k': {'pass@1': 1.0, 'pass@2': 1.0, 'pass@3': 1.0}}
+    """
+    results = evaluate_generations(generations, level=level, debug=debug)
+    metrics = get_results(results, count_errors=count_errors, k_list=k_list)
+    return metrics
+
+# import doctest
+# doctest.testmod()