update

app.py

@@ -241,15 +241,7 @@ def infilling_dream(
     x = F.pad(input_ids, (0, max_tokens - input_ids.shape[1]), value = MASK_ID)
 
     # ------ Visualization Setup
-    previous_tokens_vis = None
     initial_generated_tokens = input_ids[0, prefix_len: prefix_len + num_generation_tokens]
-    vis_data_initial = []
-    for tok_id in initial_generated_tokens.tolist():
-        display_token = MASK_TOKEN
-        color = '#4444'
-        vis_data_initial.append((display_token, color))
-
-    previous_tokens_vis = initial_generated_tokens
     #yield vis_data_initial
     yield tokenizer.decode(initial_generated_tokens.tolist()), ''
     time.sleep(visualization_delay)
@@ -330,32 +322,6 @@ def infilling_dream(
         
         ## Visualize Denoise Step
         cur_generated_tokens = x[0, prefix_len: prefix_len + num_generation_tokens]
-        print(f"Current generated tokens: {tokenizer.decode(cur_generated_tokens)}")
-        vis_data = []
-        for j in range(num_generation_tokens):
-            current_tok_id = cur_generated_tokens[j].item()
-            previous_tok_id = previous_tokens_vis[j].item() if previous_tokens_vis is not None and j < len(previous_tokens_vis) else MASK_ID
-            try:
-                decoded_token = tokenizer.decode([current_tok_id], skip_special_tokens=False)
-                display_token = MASK_TOKEN if current_tok_id == MASK_ID else decoded_token
-            except Exception: display_token = f"[ID:{current_tok_id}]"
-        vis_data = []
-        #print(f"Step {i+1}/{steps}")
-        #print(f"Current generated tokens: {tokenizer.decode(cur_generated_tokens)}")
-        # [Visualization formatting logic remains the same]
-        for j in range(num_generation_tokens):
-            current_tok_id = cur_generated_tokens[j].item()
-            previous_tok_id = previous_tokens_vis[j].item() if previous_tokens_vis is not None and j < len(previous_tokens_vis) else MASK_ID
-            try:
-                decoded_token = tokenizer.decode([current_tok_id], skip_special_tokens=False)
-                display_token = MASK_TOKEN if current_tok_id == MASK_ID else decoded_token
-            except Exception: display_token = f"[ID:{current_tok_id}]"
-
-            color = None; token_to_display = display_token
-            if current_tok_id == MASK_ID: color = "#444444"
-            else: color = "#6699CC"
-            
-            vis_data.append((token_to_display, color))
         cur_tokens =  tokenizer.decode(cur_generated_tokens.tolist())
         ## replace all <|endoftext|> with <|delete|>
         cur_tokens = cur_tokens.replace("<|endoftext|>", "<|delete|>")
@@ -408,26 +374,12 @@ def infilling_dream(
                 num_generation_tokens -= 1
 
             cur_generated_tokens = x[0, prefix_len: prefix_len + num_generation_tokens]
-            #vis_data = []
-            #for j in range(num_generation_tokens):
-            #    current_tok_id = cur_generated_tokens[j].item()
-            #    try:
-            #        decoded_token = tokenizer.decode([current_tok_id], skip_special_tokens=False)
-            #        display_token = MASK_TOKEN if current_tok_id == MASK_ID else decoded_token
-            #    except Exception: display_token = f"[ID:{current_tok_id}]"
-            #    color = None; token_to_display = display_token
-            #    if current_tok_id == MASK_ID: color = "#444444"
-            #    else: color = "#6699CC"
-
-            #    vis_data.append((token_to_display, color))
             yield tokenizer.decode(cur_generated_tokens.tolist()), ''
-            #yield vis_data
             time.sleep(visualization_delay)
 
     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()))
@@ -454,8 +406,8 @@ 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!")
         gr.Markdown(
-            "[[Model Card(link TBD)](https://huggingface.co/Dream-org/DreamOn-v0-7B)] "
-            "[[Blog(link TBD)](https://hkunlp.github.io/blog/2025/dreamon/)]"
+            "[[Model Card](https://huggingface.co/Dream-org/DreamOn-v0-7B)] "
+            "[[Blog](https://hkunlp.github.io/blog/2025/dreamon/)]"
         )
 
         with gr.Column():
@@ -506,6 +458,7 @@ def create_chatbot_demo():
             with gr.Row():
                 generate_btn = gr.Button("Generate", variant="primary")
                 clear_btn = gr.Button("Clear")
+                check_btn = gr.Button("Run test case")
 
         # Generation Settings
         with gr.Accordion("Generation Settings"):
@@ -612,6 +565,12 @@ def create_chatbot_demo():
             queue=False
         )
 
+        check_btn.click(
+            fn=check_result,
+            inputs=[output_vis, task_id_input],
+            outputs=[result_output]
+        )
+
         sample_btn.click(
             fn=get_example_input,
             outputs=[prefix_input, suffix_input, test_case_input, task_id_input],

eval.py

@@ -10,8 +10,34 @@ import tempfile
 from typing import Callable, Dict, Optional
 
 
+import subprocess
+import time
+from multiprocessing import Process, Queue
+
+def write_code_to_file(code, filename):
+    with open(filename, 'w') as f:
+        f.write(code)
+
+def read_result_from_file(filename):
+    with open(filename, 'r') as f:
+        return f.read().strip()
+
+def child_process(script_path, result_path, timeout):
+    try:
+        # Construct the command to run the script and redirect output to the result file
+        command = f'python3 {script_path} > {result_path} 2>&1'
+        subprocess.run(command, shell=True, timeout=timeout)
+    except subprocess.TimeoutExpired:
+        result = "timed out"
+    else:
+        result = read_result_from_file(result_path)
+    finally:
+        # Write the result to a file and exit
+        with open(result_path, 'w') as f:
+            f.write(result)
+
 def unsafe_execute(problem, completion, timeout=3):
-    with create_tempdir():
+    with tempfile.TemporaryDirectory() as temp_dir:
         # Construct the check program
         check_program = (
             problem["prompt"]
@@ -23,30 +49,27 @@ def unsafe_execute(problem, completion, timeout=3):
             + 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()
+        # Write the check program to a temporary file
+        script_path = os.path.join(temp_dir, 'script.py')
+        write_code_to_file(check_program, script_path)
+
+        # Create a temporary file to store the result
+        result_path = os.path.join(temp_dir, 'result.txt')
+
+        # Create a queue to communicate between parent and child processes
+        result_queue = Queue()
+
+        # Create a child process
+        p = Process(target=child_process, args=(script_path, result_path, timeout))
+        p.start()
+        p.join(timeout)
+
+        if p.is_alive():
+            p.terminate()
+            p.join()
             result = "timed out"
         else:
-            result = result_queue.get()
+            result = read_result_from_file(result_path)
 
         return result