update testcase

app.py

@@ -13,7 +13,7 @@ import traceback # For better error printing
 import random
 import gzip
 import json
-
+from eval import unsafe_execute ## args:(problem, completion)
 
 def read_problems() -> Dict[str, Dict]:
     benchmark_file = "HumanEval-SingleLineInfilling.jsonl.gz"
@@ -223,7 +223,8 @@ def infilling_dream(
     alg: str,
     alg_temp: Optional[float],
     visualization_delay: float,
-    delete_righthand_eos: bool
+    delete_righthand_eos: bool,
+    task_id: str
 ) -> List[Tuple[str, str]]:
     # ------1. Prepare the input for infilling -----------------
     prefix = prefix
@@ -249,7 +250,7 @@ def infilling_dream(
 
     previous_tokens_vis = initial_generated_tokens
     #yield vis_data_initial
-    yield tokenizer.decode(initial_generated_tokens.tolist())
+    yield tokenizer.decode(initial_generated_tokens.tolist()), ''
     time.sleep(visualization_delay)
 
     # ----2. Step by Step Infilling ----------------------------------------
@@ -357,7 +358,7 @@ def infilling_dream(
         cur_tokens =  tokenizer.decode(cur_generated_tokens.tolist())
         ## replace all <|endoftext|> with <|delete|>
         cur_tokens = cur_tokens.replace("<|endoftext|>", "<|delete|>")
-        yield cur_tokens
+        yield cur_tokens, ''
         time.sleep(visualization_delay)
 
         #  Expansion Step: Check for expand_id and replace with two mask tokens
@@ -389,7 +390,7 @@ def infilling_dream(
                 else: color = "#6699CC"
 
                 if token_to_display: vis_data.append((token_to_display, color))
-            yield tokenizer.decode(cur_generated_tokens.tolist())
+            yield tokenizer.decode(cur_generated_tokens.tolist()), ''
             #yield vis_data
             time.sleep(visualization_delay)
         ## detele EOS tokens from middle
@@ -418,23 +419,36 @@ def infilling_dream(
             #    else: color = "#6699CC"
 
             #    vis_data.append((token_to_display, color))
-            yield tokenizer.decode(cur_generated_tokens.tolist())
+            yield tokenizer.decode(cur_generated_tokens.tolist()), ''
             #yield vis_data
             time.sleep(visualization_delay)
 
-    yield tokenizer.decode(x[0, prefix_len: prefix_len + num_generation_tokens].tolist())
+    generated_code = tokenizer.decode(x[0, prefix_len: prefix_len + num_generation_tokens].tolist())
+    yield generated_code, ''
+    result = check_result(task_id, generated_code)
+    yield generated_code, result
 def get_example_input():
     ### this functions samples a case from humaneval-infilling as prefix and suffix
     task_id = random.choice(list(problems.keys()))
-    example = problems[task_id]
-    prefix, suffix = example['prompt'], example['suffix']
-    return prefix, suffix, ''
+    problem = problems[task_id]
+    prefix, suffix = problem['prompt'], problem['suffix']
+    test_case = problem['test']
+    return prefix, suffix, test_case, task_id
+
+def check_result(task_id, completion):
+    # 从数据集中读取问题
+    problem = problems[task_id]
+    # 这里假设 `unsafe_execute` 是一个可以执行代码并返回结果的函数
+    result = unsafe_execute(problem, completion)
+    return result
+
 
 
 # --- Gradio UI ---
 css = '''
 .category-legend{display:none}
 '''
+
 def create_chatbot_demo():
     with gr.Blocks(css=css) as demo:
         gr.Markdown("# DreamOn: Diffusion Language Models For Code Infilling Beyond Fixed-size Canvas\nClick **Example Prompt** to get a prefix and suffix, then click **Generate** to generate code. Have fun!")
@@ -464,6 +478,27 @@ def create_chatbot_demo():
                 lines=2
             )
 
+            # Test Case input
+            test_case_input = gr.Textbox(
+                label="Test Case",
+                placeholder="Enter your test case here...",
+                lines=2
+            )
+
+            # Hidden Task ID input
+            task_id_input = gr.Textbox(
+                label="Task ID",
+                placeholder="Task ID will be stored here...",
+                visible=False
+            )
+
+            # Result of execution
+            result_output = gr.Textbox(
+                label="Result of Execution",
+                placeholder="Execution result will be shown here...",
+                lines=2
+            )
+
             with gr.Row():
                 sample_btn = gr.Button("Example Prompt")
 
@@ -512,12 +547,6 @@ def create_chatbot_demo():
                     label="Top-P (0 disables)"
                 )
             with gr.Row():
-                top_p = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    value=0.95,
-                    step=0.05,
-                    label="Top-P (0 disables)")
                 top_k = gr.Slider(
                     minimum=0, 
                     maximum=200, 
@@ -525,7 +554,6 @@ def create_chatbot_demo():
                     step=5, 
                     label="Top-K (0 disables)")
 
-            
             with gr.Row():
                 alg = gr.Radio(
                     choices=['maskgit_plus', 'topk_margin', 'entropy'],
@@ -552,7 +580,6 @@ def create_chatbot_demo():
                     value=True
                 )
 
-
         # Connect the UI elements
         generation_inputs = [
             prefix_input,
@@ -566,71 +593,35 @@ def create_chatbot_demo():
             alg,
             alg_temp,
             visualization_delay,
-            pad_delete_righthand
+            pad_delete_righthand,
+            task_id_input
         ]
 
         generate_btn.click(
             fn=infilling_dream,
             inputs=generation_inputs,
-            outputs=[output_vis],
+            outputs=[output_vis, result_output],
             show_progress="hidden"
         )
 
         clear_btn.click(
-            lambda: ("", "", ""),  # Clear all inputs and outputs
+            lambda: ("", "", "", "", "", ""),  # Clear all inputs and outputs
             inputs=[],
-            outputs=[prefix_input, suffix_input, output_vis],
+            outputs=[prefix_input, suffix_input, output_vis, test_case_input, result_output, task_id_input],
             queue=False
         )
 
         sample_btn.click(
             fn=get_example_input,
-            outputs=[prefix_input, suffix_input, output_vis],
+            outputs=[prefix_input, suffix_input, test_case_input, task_id_input],
             queue=False
         )
     return demo
 
-def test():
-    prefix = '''import List
 
-def has_close_elements(numbers: List[float], threshold: float) -> bool:
-    """ Check if in given list of numbers, are any two numbers closer to each other than
-    given threshold.
-    >>> has_close_elements([1.0, 2.0, 3.0], 0.5)
-    False
-    >>> has_close_elements([1.0, 2.8, 3.0, 4.0, 5.0, 2.0], 0.3)
-    True
-    """
-
-'''
-
-    suffix = '''
-        for idx2, elem2 in enumerate(numbers):
-            if idx != idx2:
-                distance = abs(elem - elem2)
-                if distance < threshold:
-                    return True
-
-    return False
-'''
-    infilling_dream(
-    prefix=prefix,
-    suffix= suffix,
-    start_gen_len=4,
-    max_gen_len= 64,
-    expand_budget= 60,
-    temperature=0.2,
-    top_p = 0.9,
-    top_k = None,
-    alg= 'entropy',
-    alg_temp=0,
-    visualization_delay= 0,
-    delete_righthand_eos= True
-)
 
 # --- Launch ---
 if __name__ == "__main__":
     #test()
     demo = create_chatbot_demo()
-    demo.queue().launch(debug=True, share=True)
-    
+    demo.queue().launch(debug=True, share=True)

eval.py

@@ -0,0 +1,201 @@
+import ast
+import contextlib
+import faulthandler
+import io
+import multiprocessing
+import os
+import platform
+import signal
+import tempfile
+from typing import Callable, Dict, Optional
+
+
+def unsafe_execute(problem, completion, timeout=3):
+    with create_tempdir():
+        # Construct the check program
+        check_program = (
+            problem["prompt"]
+            + completion
+            + problem["suffix"]
+            + "\n"
+            + problem["test"]
+            + "\n"
+            + f"check({problem['entry_point']})"
+        )
+
+        # Use multiprocessing to execute the code in a separate process
+        result_queue = multiprocessing.Queue()
+
+        def worker(check_program, result_queue):
+            try:
+                exec_globals = {}
+                with swallow_io():
+                    with time_limit(timeout):
+                        exec(check_program, exec_globals)
+                result_queue.put("passed")
+            except TimeoutException:
+                result_queue.put("timed out")
+            except BaseException as e:
+                result_queue.put(f"failed: {e}")
+
+        process = multiprocessing.Process(target=worker, args=(check_program, result_queue))
+        process.start()
+        process.join(timeout + 1)  # Give some extra time for cleanup
+
+        if process.is_alive():
+            process.terminate()
+            result = "timed out"
+        else:
+            result = result_queue.get()
+
+        return result
+
+
+@contextlib.contextmanager
+def time_limit(seconds: float):
+    def signal_handler(signum, frame):
+        raise TimeoutException("Timed out!")
+
+    signal.setitimer(signal.ITIMER_REAL, seconds)
+    signal.signal(signal.SIGALRM, signal_handler)
+    try:
+        yield
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, 0)
+
+
+@contextlib.contextmanager
+def swallow_io():
+    stream = WriteOnlyStringIO()
+    with contextlib.redirect_stdout(stream):
+        with contextlib.redirect_stderr(stream):
+            with redirect_stdin(stream):
+                yield
+
+
+@contextlib.contextmanager
+def create_tempdir():
+    with tempfile.TemporaryDirectory() as dirname:
+        with chdir(dirname):
+            yield dirname
+
+
+class TimeoutException(Exception):
+    pass
+
+
+class WriteOnlyStringIO(io.StringIO):
+    """StringIO that throws an exception when it's read from"""
+
+    def read(self, *args, **kwargs):
+        raise IOError
+
+    def readline(self, *args, **kwargs):
+        raise IOError
+
+    def readlines(self, *args, **kwargs):
+        raise IOError
+
+    def readable(self, *args, **kwargs):
+        """Returns True if the IO object can be read."""
+        return False
+
+
+class redirect_stdin(contextlib._RedirectStream):  # type: ignore
+    _stream = "stdin"
+
+
+@contextlib.contextmanager
+def chdir(root):
+    if root == ".":
+        yield
+        return
+    cwd = os.getcwd()
+    os.chdir(root)
+    try:
+        yield
+    except BaseException as exc:
+        raise exc
+    finally:
+        os.chdir(cwd)
+
+
+def reliability_guard(maximum_memory_bytes: Optional[int] = 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