Upload 7 files

README.md

@@ -1,14 +1,15 @@
 ---
-title: MGDebugger
-emoji: 😻
-colorFrom: blue
-colorTo: pink
+title: Test
+emoji: 📈
+colorFrom: pink
+colorTo: green
 sdk: gradio
 sdk_version: 4.44.1
-app_file: app.py
+app_file: space_demo.py
 pinned: false
-license: apache-2.0
-short_description: 'MGDebugger Demo: Multi-Granularity LLM Debugger'
+license: mit
 ---
 
+We are highly inspired by [LDB](https://huggingface.co/spaces/shangdatalab-ucsd/LDB) 🚀✨
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

main.py

@@ -0,0 +1,640 @@
+import time
+from openai import OpenAI
+from collections import Counter
+from loguru import logger
+import json
+import os
+from tqdm import tqdm
+import sys
+import io
+import traceback
+import ast
+import time
+import re
+from groq import Groq
+from utils import split_nested_functions, get_dependency_graph_str, evaluate, parse_json_response, extract_code_blocks, extract_functions, extract_function, create_dependency_graph, topological_sort, merge_changes_to_parents, evaluate_simple, parse_transcoder_problem_content
+from testcase_utils import get_parameter_names, parse_tests
+
+# CUDA_VISIBLE_DEVICES=0 python -m vllm.entrypoints.openai.api_server --model deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct --trust-remote-code --dtype auto --api-key token-abc123s --port 18889
+# CUDA_VISIBLE_DEVICES=0 python -m vllm.entrypoints.openai.api_server --model TechxGenus/Codestral-22B-v0.1-GPTQ --dtype auto --api-key token-abc123s --port 18890 --trust-remote-code --chat-template helper/codestral_template.jinja
+# CUDA_VISIBLE_DEVICES=0 python -m vllm.entrypoints.openai.api_server --model Qwen/CodeQwen1.5-7B-Chat --dtype auto --api-key token-abc123s --port 18892 --trust-remote-code
+
+MODEL = "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct"
+# MODEL = "TechxGenus/Codestral-22B-v0.1-GPTQ"
+# MODEL = "Qwen/CodeQwen1.5-7B-Chat"
+
+
+client = OpenAI(
+    base_url="http://localhost:18889/v1",
+    api_key="token-abc123s",
+)
+
+# dscoder 18889
+# codestral 18890
+# codeqwen 18892
+
+# Hyperparameters
+MAX_VLLM_RETRIES = 10  # maximum number of retries for the VLLM call
+MAX_PARSE_RETRIES = 3  # we sometimes fail to parse the code/test cases from the response, so we retry
+MAX_DEBUG_RETRIES = 1  # 3 seems to be slightly better than 1, but the difference is not significant
+RETRY_DELAY = 0  # seconds
+REPEAT_CONVERT_HIERARCHICAL_NUM = 1  # seems unimportant
+REPEAT_TEST_CASE_GENERATION_NUM = 1  # 1 seems to be better than 3
+MAX_OUTER_RETRY = 10  # maximum number of retries for the entire debugging process
+CONTINUOUS_RETRY = False  # whether to retry from the last fixed code, else retry from the original buggy code. it seems better to set this to False
+TEMPERATURE = 0.8  # 0.8 better than 1.0 better than 0.2
+MINP = 0.05
+
+# for collecting stats
+TOTAL_PROMPT_TOKENS = 0
+TOTAL_COMPLETION_TOKENS = 0
+TOTAL_MG_DEBUG_CALLS = 0
+TOTAL_DEBUG_FUNCTION_CALLS = 0
+TOTAL_GENERATE_TEST_CASES_CALLS = 0
+TOTAL_CONVERT_HIERARCHICAL_CALLS = 0
+
+
+def get_completion_with_retry(messages, model=MODEL, MAX_VLLM_RETRIES=MAX_VLLM_RETRIES):
+    global TOTAL_PROMPT_TOKENS, TOTAL_COMPLETION_TOKENS
+    for attempt in range(MAX_VLLM_RETRIES):
+        try:
+            logger.info(f"Attempting LLM call (attempt {attempt + 1}/{MAX_VLLM_RETRIES})")
+            logger.info(f"Input messages: {messages[-1]['content']}")
+            chat_completion = client.chat.completions.create(
+                messages=messages,
+                model=MODEL,
+                temperature=TEMPERATURE
+            )
+            response = chat_completion.choices[0].message.content
+            logger.info(f"LLM response: {response}")
+            logger.info("LLM call received")
+
+            # Update token counts
+            TOTAL_PROMPT_TOKENS += chat_completion.usage.prompt_tokens
+            TOTAL_COMPLETION_TOKENS += chat_completion.usage.completion_tokens
+
+            return response
+        except Exception as e:
+            logger.error(f"LLM call failed: {str(e)}")
+            if attempt < MAX_VLLM_RETRIES - 1:
+                logger.info(f"Retrying in {RETRY_DELAY} seconds...")
+                time.sleep(RETRY_DELAY)
+            else:
+                logger.error("Max retries reached. Giving up.")
+                raise
+
+
+def generate_test_cases(full_code, gold_test_cases, function_name, MAX_PARSE_RETRIES=MAX_PARSE_RETRIES):
+    global TOTAL_GENERATE_TEST_CASES_CALLS
+    TOTAL_GENERATE_TEST_CASES_CALLS += 1
+
+    logger.info(f"Generating test cases for function: {function_name}")
+
+    prompt = f"""
+    Analyze the following Python code and focus on the function named `{function_name}`. 
+    Generate the same number of test cases for this specific function based on the provided gold test cases for the main function.
+    Ensure that the generated test cases are consistent with the behavior expected in the gold test cases.
+
+    ### Full Code
+    {full_code}
+
+    ### Gold Test Cases for the main function
+    {json.dumps(gold_test_cases, indent=2)}
+
+    ### Function to generate test cases for: `{function_name}`
+
+    ### Output format:
+    **Test Case 1:**
+    Input: ...
+    Analysis: ...
+    Expected Output: ...
+    
+    **Test Case 2:**
+    ...
+    
+    **All Test Cases:**
+    ```json
+    {{
+        "test_cases": [
+            {{"input": {{'var1': 'value1', 'var2': 'value2'}}, "expected_output": "expected_output"}},
+            ...
+        ]
+    }}
+    ```
+    
+    ### Hint
+    Analyze how the `{function_name}` function is used within the main function and how it contributes to the expected outputs in the gold test cases. Generate test cases that reflect this behavior. For each test case, you should analyze step-by-step based on both the input and the expected output of the main function, and then provide the corresponding input and expected output for the `{function_name}` function.
+    """
+
+
+    messages = [
+        {'role': 'system', 'content': "You are an AI assistant specialized in analyzing Python functions and generating test cases."},
+        {'role': 'user', 'content': prompt},
+    ]
+
+    if "starcoder2" in MODEL.lower():
+        # remove the system message
+        messages = messages[1:]
+
+    all_test_cases = []
+    for _ in range(REPEAT_TEST_CASE_GENERATION_NUM):
+        response = get_completion_with_retry(messages)
+        parsed_response = parse_json_response(response)
+        if parsed_response and 'test_cases' in parsed_response:
+            all_test_cases.extend(parsed_response['test_cases'])
+
+    if not all_test_cases:
+        logger.error("Failed to generate any valid test cases")
+        return []
+
+    # Perform majority voting on expected outputs
+    logger.info(f"Generated {len(all_test_cases)} test cases for {function_name}: {all_test_cases}")
+    test_case_counter = Counter([json.dumps(tc) for tc in all_test_cases])
+    # show the frequency of each test case
+    logger.info(f"Test case frequency:")
+    for tc, freq in test_case_counter.items():
+        logger.info(f"Test case: {tc}, Frequency: {freq}")
+    most_common_test_cases = test_case_counter.most_common(len(gold_test_cases))
+    final_test_cases = [json.loads(tc[0]) for tc in most_common_test_cases]
+
+    logger.info(f"Voted {len(final_test_cases)} test cases for {function_name}: {final_test_cases}")
+    return final_test_cases
+
+
+def debug_function(function_code, function_name, test_cases, MAX_PARSE_RETRIES=MAX_PARSE_RETRIES):
+    global TOTAL_DEBUG_FUNCTION_CALLS
+    TOTAL_DEBUG_FUNCTION_CALLS += 1
+
+    logger.info(f"Debugging function:\n{function_code}\nWith test cases: {test_cases}")
+    prompt = f"""
+    Debug the following Python function. The function is not passing all test cases.
+    Analyze the code, identify the bug, and provide a fixed version of the function.
+
+    ### Function:
+    {function_code}
+
+    ### Test Cases:
+    {json.dumps(test_cases, indent=2)}
+    
+    ### Task:
+
+    Provide your response in the following format:
+    1. try to work as a python interpreter to execute the code step-by-step.
+    2. identify the change of each variable as you "run" the code line-by-line.
+    3. based on the execution trace, try to identify the bug
+    4. provide the final fixed code in a Python code block (```python ... ```)
+
+    Make sure to include the entire function, including the function signature. And you will be rewarded to simulate the code execution in your mind and provide step-by-step trace of the code execution.
+    
+    """
+
+    messages = [
+        {'role': 'system', 'content': 'You are an AI assistant helping to debug Python functions.'},
+        {'role': 'user', 'content': prompt},
+    ]
+
+    if "starcoder2" in MODEL.lower():
+        # remove the system message
+        messages = messages[1:]
+
+    for attempt in range(MAX_PARSE_RETRIES):
+        try:
+            response = get_completion_with_retry(messages)
+            code_blocks = extract_code_blocks(response)
+            if code_blocks:
+                # search in a reverse order to find the final fixed function
+                for block in code_blocks[::-1]:
+                    fixed_function = extract_function(block, function_name)
+                    if fixed_function:
+                        analysis = response.split("```python")[0].strip()
+                        logger.info("Generated debug analysis and fixed code")
+                        return analysis, fixed_function
+            raise ValueError("No valid fixed function found in the response")
+        except ValueError as e:
+            logger.error(f"Failed to extract fixed function (attempt {attempt + 1}/{MAX_PARSE_RETRIES}): {str(e)}")
+            if attempt < MAX_PARSE_RETRIES - 1:
+                logger.info(f"Retrying in {RETRY_DELAY} seconds...")
+                time.sleep(RETRY_DELAY)
+            else:
+                logger.error("Max retries reached. Returning None for analysis and fixed code.")
+                return None, None
+
+
+def convert_to_hierarchical(code, include_example=False):
+    global TOTAL_CONVERT_HIERARCHICAL_CALLS
+    TOTAL_CONVERT_HIERARCHICAL_CALLS += 1
+
+    logger.info("Converting code to hierarchical structure")
+
+    example = """
+    ### Example of tree-style hierarchical structure:
+    
+    ```python
+    def main_function(input):
+        preprocessed_data = preprocess(input)
+        result = process(preprocessed_data)
+        return result
+
+    def preprocess(data):
+        cleaned_data = clean_data(data)
+        normalized_data = normalize_data(cleaned_data)
+        return normalized_data
+
+    def clean_data(data):
+        # Implementation of data cleaning
+        pass
+
+    def normalize_data(data):
+        # Implementation of data normalization
+        pass
+
+    def process(data):
+        feature_vector = extract_features(data)
+        result = classify(feature_vector)
+        return result
+
+    def extract_features(data):
+        # Implementation of feature extraction
+        pass
+
+    def classify(feature_vector):
+        # Implementation of classification
+        pass
+    ```
+    """ if include_example else ""
+
+    prompt = f"""
+    Convert the following Python code into a tree-style hierarchical structure with multiple levels of sub-functions.
+    Each significant step or logical block should be its own function, and functions can call other sub-functions.
+    Ensure that the main function calls these sub-functions in the correct order, creating a tree-like structure.
+
+    ### Original Code:
+    {code}
+
+    {example}
+
+    ### Instructions:
+    Please first analyze the codes step by step, and then provide the converted code in a Python code block (```python ... ```). When providing the final converted code, make sure to include all the functions in a flattened format, where each function is defined separately.
+    """
+
+    messages = [
+        {'role': 'system', 'content': 'You are an AI assistant specialized in refactoring Python code into a tree-style hierarchical structure.'},
+        {'role': 'user', 'content': prompt},
+    ]
+
+    if "starcoder2" in MODEL.lower():
+        
+        # remove the system message
+        messages = messages[1:]
+
+    best_conversion = None
+    max_subfunctions = 0
+
+    for _ in range(REPEAT_CONVERT_HIERARCHICAL_NUM):
+        response = get_completion_with_retry(messages)
+        code_blocks = extract_code_blocks(response)
+
+        if code_blocks:
+            converted_code = code_blocks[0]
+            subfunctions = len(extract_functions(converted_code)) - 1  # Subtract 1 to exclude the main function
+
+            if subfunctions > max_subfunctions:
+                max_subfunctions = subfunctions
+                best_conversion = converted_code
+
+    if best_conversion:
+        logger.info(f"Converted code to tree-style hierarchical structure with {max_subfunctions} sub-functions")
+        # split nested functions
+        best_conversion = split_nested_functions(best_conversion)
+        return best_conversion
+    else:
+        logger.error("Failed to convert code to tree-style hierarchical structure")
+        code = split_nested_functions(code)
+        return code
+
+def mg_debug_demo(full_code, gold_test_cases, max_debug_attempts=MAX_DEBUG_RETRIES, given_model=MODEL, given_client=client):
+    global MODEL, client
+    MODEL = given_model
+    client = given_client
+    return mg_debug(full_code, gold_test_cases, max_debug_attempts)
+
+
+def mg_debug(full_code, gold_test_cases, max_debug_attempts=MAX_DEBUG_RETRIES):
+    global TOTAL_MG_DEBUG_CALLS
+    TOTAL_MG_DEBUG_CALLS += 1
+    logger.info("Starting main debugging process")
+
+    convert_hierarchical_attempts = 0
+    while convert_hierarchical_attempts < MAX_PARSE_RETRIES:
+        try:
+            # Convert to tree-style hierarchical structure
+            hierarchical_code = convert_to_hierarchical(full_code, include_example=False)
+            logger.info(f"Converted code to tree-style hierarchical structure:\n{hierarchical_code}")
+
+            functions = extract_functions(hierarchical_code)
+
+            # Create a dependency graph
+            dependency_graph = create_dependency_graph(functions)
+            logger.info(f"Dependency graph:\n{get_dependency_graph_str(dependency_graph)}")
+            
+            break
+        except Exception as e:
+            logger.error(f"Failed to convert code to hierarchical structure (attempt {convert_hierarchical_attempts + 1}/{MAX_PARSE_RETRIES}): {str(e)}")
+            convert_hierarchical_attempts += 1
+            # retry
+
+    # Sort functions based on their dependencies (bottom-up)
+    sorted_functions = topological_sort(dependency_graph)
+    logger.info(f"Sorted functions: {sorted_functions}")
+
+    for func_name in sorted_functions:
+        logger.info(f"Processing function: {func_name}")
+
+        func_code = functions[func_name]
+        test_cases = generate_test_cases(hierarchical_code, gold_test_cases, func_name)
+        fixed_code = func_code
+
+        for debug_attempt in range(max_debug_attempts):
+            all_tests_pass = True
+
+            for test_case in test_cases:
+                passed, result = evaluate(hierarchical_code, func_name, test_case)
+                if not passed:
+                    all_tests_pass = False
+                    break
+
+            if all_tests_pass:
+                logger.info(f"All tests passed for function: {func_name}")
+                break
+
+            logger.info(f"Debugging function: {func_name} (Attempt {debug_attempt + 1}/{max_debug_attempts})")
+            analysis, new_fixed_code = debug_function(fixed_code, func_name, test_cases)
+
+            if new_fixed_code:
+                logger.info(f"New fixed code for {func_name}:\n{new_fixed_code}")
+                fixed_code = new_fixed_code
+                functions[func_name] = fixed_code
+                logger.info(f"Merging {func_name} changes")
+                hierarchical_code = merge_changes_to_parents(func_name, dependency_graph, functions)
+                logger.info(f"Code after merging updates in {func_name}:\n{hierarchical_code}")
+            else:
+                logger.warning(f"Failed to fix {func_name}. Keeping previous implementation.")
+                break
+
+        if not all_tests_pass:
+            logger.warning(f"Could not fix {func_name} after {max_debug_attempts} attempts. Keeping original implementation.")
+
+    # Reconstruct the full code with fixed functions
+    fixed_full_code = "\n\n".join(functions.values())
+    logger.info("Debugging process completed. Reconstructed full code.")
+
+    return fixed_full_code
+
+
+def test():
+    buggy_code = '''
+def make_palindrome(string: str) -> str:
+    """ Find the shortest palindrome that begins with the supplied string. """
+    
+    def is_palindrome(s: str) -> bool:
+        return s == s[::-1]
+
+    suffix_start = 0
+    for i in range(len(string)):
+        if is_palindrome(string[i:]):
+            suffix_start = i
+            
+    return string + string[:suffix_start][::-1]
+    '''.strip()
+
+    gold_test_cases = [
+        {'input': 'cat', 'expected_output': 'catac'},
+        {'input': 'cata', 'expected_output': 'catac'},
+        {'input': '', 'expected_output': ''}
+    ]
+
+    entry_point = 'make_palindrome'
+
+    fixed_code = mg_debug(buggy_code, gold_test_cases)
+    logger.info(f"Fixed code:\n{fixed_code}")
+
+    logger.info("============= Final evaluation with private test cases =============")
+    # evaluate the final codes with private testcases
+    all_tests = gold_test_cases
+    for testcase in all_tests:
+        result, testcase = evaluate(fixed_code, entry_point, testcase)
+        logger.info(f"Passed: {result}, Test case: {testcase}")
+
+
+def debug_humaneval(input_seeds: str, max_examples: int = None, output_folder: str = None):
+    fixed_problems = 0
+    total_unsolved = 0
+
+    with open(input_seeds, "r") as f:
+        seeds = f.readlines()
+
+    unsolved_seeds = []
+    # filter those problems that are not solved
+    for i in range(len(seeds)):
+        problem = json.loads(seeds[i])
+        if not problem["is_solved"]:
+            unsolved_seeds.append(problem)
+
+    unsolved_seeds = [seed for seed in unsolved_seeds if seed['task_id'] == 'HumanEval/38']
+    
+    # resume_path = "output_data/mbpp/seed/codeqwen/20240909-025624/CodeQwen1.5-7B-Chat_debugging_seeds_from_codeqwen.jsonl"
+    # # only load "debugged": false problems from the resume path
+    # with open(resume_path, "r") as f:
+    #     resume_problems = f.readlines()
+    # resume_unsolved_task_ids = [json.loads(problem)["task_id"] for problem in resume_problems if not json.loads(problem)["debugged"]]
+    # unsolved_seeds = [problem for problem in unsolved_seeds if problem["task_id"] in resume_unsolved_task_ids]
+    # logger.info(f"Resuming from {resume_path}, {len(unsolved_seeds)} problems to debug")
+
+    # filter the unsolved problems that are not in the resume path
+
+    # parse transcoder problems
+    if "transcoder" in input_seeds.lower():
+        logger.info(f"Parsing the problem content for transcoder problems")
+        unsolved_seeds = [parse_transcoder_problem_content(problem) for problem in tqdm(unsolved_seeds)]
+
+    total_unsolved = len(unsolved_seeds)
+    logger.info(f"Debugging {total_unsolved} unsolved problems")
+    if max_examples is not None:
+        unsolved_seeds = unsolved_seeds[:max_examples]
+        logger.info(f"Filtering to {max_examples} examples")
+
+    for problem in tqdm(unsolved_seeds, ncols=100):
+        
+        model_to_be_fixed = input_seeds.split("/")[-2]
+        model_name = MODEL.split("/")[-1]
+        with open(f"{output_folder}/{model_name}_debugging_seeds_from_{model_to_be_fixed}.jsonl", "w+") as f:
+            for seed in unsolved_seeds:
+                f.write(json.dumps(seed) + "\n")
+                
+        logger.info(f"Processing unsolved problem: {problem['task_id']}")
+        logger.info(f"Problem: {problem}")
+        logger.info(f"Problem Raw Prompt: \n{problem['prompt']}")
+
+        try:
+            buggy_code = problem["solution"]
+            entry_point = problem["entry_point"]
+            try:
+                parameter_names = get_parameter_names(problem["prompt"], entry_point)
+            except:
+                parameter_names = get_parameter_names(problem["solution"], entry_point)
+            logger.info(f"Extracted parameter names: {parameter_names}")
+
+            # in order to save time, we extract the first 3 given tests for transcoder
+            if "transcoder" in problem["task_id"].lower():
+                logger.info(f"Extracted {len(problem['given_tests'])} given tests, only using the first 3 samples")
+                problem["given_tests"] = problem["given_tests"][:3]
+
+            gold_tests_raw = "\n".join(problem["given_tests"]).replace(entry_point, "candidate")
+            gold_tests = parse_tests(gold_tests_raw, parameter_names, entry_point)["test_cases"]
+            logger.info(f"Extracted gold test cases: {gold_tests}")
+
+            problem['fixed_codes'] = []  # Initialize list to store fixed codes for each retry
+            problem['mg_debug_retries'] = 0  # Initialize retry counter
+
+            for outer_retry in range(MAX_OUTER_RETRY):
+                fixed_code = mg_debug(buggy_code, gold_tests)
+                problem['fixed_codes'].append(fixed_code)
+                problem['mg_debug_retries'] += 1
+
+                all_passed = evaluate_simple(fixed_code, entry_point, problem["test"])
+                if all_passed:
+                    fixed_problems += 1
+                    problem['debugged'] = True
+                    logger.info(f"Successfully fixed problem: {problem['task_id']} on retry {outer_retry + 1}")
+                    break
+                else:
+                    logger.info(f"Failed to fix problem: {problem['task_id']} on retry {outer_retry + 1}")
+                    if outer_retry < MAX_OUTER_RETRY - 1:
+                        if CONTINUOUS_RETRY:
+                            buggy_code = fixed_code  # Use the last fixed code as the new buggy code for the next retry
+                        else:
+                            # Use the original buggy code for the next retry
+                            buggy_code = problem["solution"]
+
+            if not all_passed:
+                problem['debugged'] = False
+                logger.info(f"Failed to fix problem: {problem['task_id']} after {MAX_OUTER_RETRY} retries")
+
+        except Exception as e:
+            logger.error(f"Error occurred while processing problem: {problem['task_id']}")
+            logger.error(traceback.format_exc())
+            problem['fixed_codes'] = []
+            problem['debugged'] = False
+            problem['mg_debug_retries'] = 0
+
+    # Add statistics to the log
+    logger.info("=== Statistics ===")
+    logger.info(f"Total prompt tokens: {TOTAL_PROMPT_TOKENS}")
+    logger.info(f"Total completion tokens: {TOTAL_COMPLETION_TOKENS}")
+    logger.info(f"Total mg_debug calls: {TOTAL_MG_DEBUG_CALLS}")
+    logger.info(f"Total debug_function calls: {TOTAL_DEBUG_FUNCTION_CALLS}")
+    logger.info(f"Total generate_test_cases calls: {TOTAL_GENERATE_TEST_CASES_CALLS}")
+    logger.info(f"Total convert_hierarchical calls: {TOTAL_CONVERT_HIERARCHICAL_CALLS}")
+
+    # Compute and log average statistics
+    avg_prompt_tokens = TOTAL_PROMPT_TOKENS / TOTAL_MG_DEBUG_CALLS if TOTAL_MG_DEBUG_CALLS > 0 else 0
+    avg_completion_tokens = TOTAL_COMPLETION_TOKENS / TOTAL_MG_DEBUG_CALLS if TOTAL_MG_DEBUG_CALLS > 0 else 0
+    avg_debug_function_calls = TOTAL_DEBUG_FUNCTION_CALLS / TOTAL_MG_DEBUG_CALLS if TOTAL_MG_DEBUG_CALLS > 0 else 0
+    avg_generate_test_cases_calls = TOTAL_GENERATE_TEST_CASES_CALLS / TOTAL_MG_DEBUG_CALLS if TOTAL_MG_DEBUG_CALLS > 0 else 0
+
+    logger.info(f"Average prompt tokens per mg_debug call: {avg_prompt_tokens:.2f}")
+    logger.info(f"Average completion tokens per mg_debug call: {avg_completion_tokens:.2f}")
+    logger.info(f"Average debug_function calls per mg_debug call: {avg_debug_function_calls:.2f}")
+    logger.info(f"Average generate_test_cases calls per mg_debug call: {avg_generate_test_cases_calls:.2f}")
+
+    # distribution of debug retries that solved the problem
+    debug_retries = [problem['mg_debug_retries'] for problem in unsolved_seeds if problem['debugged']]
+    debug_retries_counter = Counter(debug_retries)
+
+    logger.info(f"=== Final Results ===")
+    logger.info(f"Total unsolved problems in seeds: {total_unsolved}")
+    logger.info(f"Problems fixed by our method: {fixed_problems}")
+    logger.info(f"Success rate: {fixed_problems / total_unsolved * 100:.2f}%")
+    logger.info(f"Debug retries distribution for solved problems: {debug_retries_counter}")
+
+    logger.info(f"Total solved problems before: {len(seeds) - total_unsolved}")
+    logger.info(f"Total solved problems after: {len(seeds) - total_unsolved + fixed_problems}")
+    logger.info(f"Previous accuracy: {(len(seeds) - total_unsolved) / len(seeds) * 100:.2f}%")
+    logger.info(f"Final accuracy: {(len(seeds) - total_unsolved + fixed_problems) / len(seeds) * 100:.2f}%")
+
+    # save the final results to a file
+    model_to_be_fixed = input_seeds.split("/")[-2]
+    model_name = MODEL.split("/")[-1]
+    with open(f"{output_folder}/{model_name}_debugging_seeds_from_{model_to_be_fixed}.jsonl", "w+") as f:
+        for seed in unsolved_seeds:
+            f.write(json.dumps(seed) + "\n")
+
+    with open(f"{output_folder}/statistics.json", "w") as f:
+        stats = {
+            "total_prompt_tokens": TOTAL_PROMPT_TOKENS,
+            "total_completion_tokens": TOTAL_COMPLETION_TOKENS,
+            "total_mg_debug_calls": TOTAL_MG_DEBUG_CALLS,
+            "total_debug_function_calls": TOTAL_DEBUG_FUNCTION_CALLS,
+            "total_generate_test_cases_calls": TOTAL_GENERATE_TEST_CASES_CALLS,
+            "total_convert_hierarchical_calls": TOTAL_CONVERT_HIERARCHICAL_CALLS,
+            "avg_prompt_tokens_per_mg_debug": avg_prompt_tokens,
+            "avg_completion_tokens_per_mg_debug": avg_completion_tokens,
+            "avg_debug_function_calls_per_mg_debug": avg_debug_function_calls,
+            "avg_generate_test_cases_calls_per_mg_debug": avg_generate_test_cases_calls
+        }
+        json.dump(stats, f, indent=2)
+        scores = {
+            "total_unsolved": total_unsolved,
+            "fixed_problems": fixed_problems,
+            "success_rate": fixed_problems / total_unsolved * 100,
+            "debug_retries_distribution": debug_retries_counter,
+            "total_solved_problems_before": len(seeds) - total_unsolved,
+            "total_solved_problems_after": len(seeds) - total_unsolved + fixed_problems,
+            "previous_accuracy": (len(seeds) - total_unsolved) / len(seeds) * 100,
+            "final_accuracy": (len(seeds) - total_unsolved + fixed_problems) / len(seeds) * 100
+        }
+        json.dump(scores, f, indent=2)
+
+
+if __name__ == "__main__":
+
+    input_seeds = "input_data/humaneval/seed/reflexion/seed.jsonl"
+    # input_seeds = "input_data/humaneval/seed/codestral/seed.jsonl"
+    # input_seeds = "input_data/humaneval/seed/codeqwen/seed.jsonl"
+
+    # input_seeds = "input_data/mbpp/seed/codestral/seed.jsonl"
+    # input_seeds = "input_data/mbpp/seed/deepseekcoder/seed.jsonl"
+    # input_seeds = "input_data/mbpp/seed/codeqwen/seed.jsonl"
+    
+    # input_seeds = "input_data/humanevalfix/seeds.jsonl"
+
+    seed_stamp = input_seeds.split("input_data/")[-1].replace("/seed.jsonl", "")
+    timestamp = time.strftime("%Y%m%d-%H%M%S")
+    output_folder = f"./output_data/{seed_stamp}/{timestamp}"
+    os.makedirs(output_folder, exist_ok=True)
+
+    # Configure logger
+    logger.remove()
+    logger.add(sys.stderr, format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>")
+    logger.add(f"{output_folder}/all_info.log", format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>")
+
+    # save all important params in a file
+    with open(f"{output_folder}/params.json", "w+") as f:
+        info = {
+            "input_seeds": input_seeds,
+            "output_folder": output_folder,
+            "MAX_DEBUG_RETRIES": MAX_DEBUG_RETRIES,
+            "MAX_PARSE_RETRIES": MAX_PARSE_RETRIES,
+            "RETRY_DELAY": RETRY_DELAY,
+            "REPEAT_CONVERT_HIERARCHICAL_NUM": REPEAT_CONVERT_HIERARCHICAL_NUM,
+            "REPEAT_TEST_CASE_GENERATION_NUM": REPEAT_TEST_CASE_GENERATION_NUM,
+            "MODEL": MODEL,
+            "MAX_OUTER_RETRY": MAX_OUTER_RETRY,
+            "CONTINUOUS_RETRY": CONTINUOUS_RETRY,
+            "TEMPERATURE": TEMPERATURE,
+            "MAX_VLLM_RETRIES": MAX_VLLM_RETRIES,
+            "MINP": MINP
+        }
+        f.write(json.dumps(info, indent=2))
+
+    debug_humaneval(input_seeds, max_examples=None, output_folder=output_folder)

requirements.txt

@@ -0,0 +1,13 @@
+jsonlines==3.1.0
+openai==1.12.0
+datasets
+tenacity==8.1.0
+astunparse
+transformers
+accelerate
+astor
+graphviz
+vllm
+astroid
+groq
+loguru

space_demo.py

@@ -0,0 +1,180 @@
+import re
+import pandas as pd
+import gradio as gr
+import ast
+import random
+from main import mg_debug_demo
+from groq import Groq
+from openai import OpenAI
+
+
+def debug_code(key, model, code, fixed_code, test_cases):
+    if model == 'deepseek-coder':
+        client = OpenAI(
+            api_key=key,
+            base_url="https://api.deepseek.com"
+        )
+    elif model in ['gemma2-9b-it']:
+        client = Groq(
+            api_key=key,
+        )
+    else:
+        client = OpenAI(
+            api_key=key
+        )
+    fixed_code = mg_debug_demo(code, test_cases, given_model=model, given_client=client)
+    return code, fixed_code, test_cases
+
+
+app = gr.Blocks(
+    theme=gr.themes.Default(primary_hue="red", secondary_hue="pink", neutral_hue="gray")
+)
+
+with app:
+    with gr.Row():
+        gr.Markdown("# MGDebugger Demo")
+    with gr.Row():
+        with gr.Column():
+            with gr.Row():
+                openai_key_input = gr.Textbox(
+                    label="API Key",
+                    placeholder="Enter your API key here",
+                    type="password",
+                )
+                model_selector = gr.Dropdown(
+                    label="Choose Model",
+                    choices=[("gemma2. Please get your api key from https://console.groq.com/keys", "gemma2-9b-it"),
+                             ("deepseek-coder. Please get your api key from https://platform.deepseek.com/api-docs/", "deepseek-coder"), 
+                             ("gpt-4o-mini. Please get your api key from https://platform.openai.com/settings/profile?tab=api-keys", "gpt-4o-mini"),
+                             ("gpt-4-1106-preview. Please get your api key from https://platform.openai.com/settings/profile?tab=api-keys", "gpt-4-1106-preview")],
+                    value="deepseek-coder",
+                )
+            code_input = gr.TextArea(
+                label="Code Input",
+                placeholder="Enter your code here",
+                lines=10,
+                )
+            test_cases = gr.TextArea(
+                label="Public Test Cases",
+                placeholder="Enter your public test cases",
+                lines=10,
+            )
+            with gr.Row():  # This Row will contain the buttons
+                debug_button = gr.Button("Debug", variant="primary")
+                clear_button = gr.Button("Clear", variant="neutral")
+        with gr.Column():
+            fixed_code_output = gr.TextArea(
+                label="Fixed Code",
+                placeholder="Fixed code will be shown here",
+                lines=10,
+                interactive=False,
+                visible=True,
+            )
+    debug_button.click(
+        debug_code,
+        inputs=[
+            openai_key_input,
+            model_selector,
+            code_input,
+            fixed_code_output,
+            test_cases,
+        ],
+        outputs=[ code_input, fixed_code_output, test_cases],
+    )
+
+    def clear_inputs():
+        return (
+            "",
+            "",
+            pd.DataFrame(
+                {
+                    "Pass?": [],
+                    "Expression": [],
+                    "Expected Value": [],
+                    "Actual Value": [],
+                }
+            ),
+            "",
+            "",
+        )
+
+    clear_button.click(
+        clear_inputs,
+        inputs=[],
+        outputs=[code_input, test_cases, fixed_code_output],
+    )
+
+    gr.Markdown("## Text Examples")
+    gr.Examples(
+        [
+            [
+                
+                {'input': {'music_string': 'o o| .| o| o| .| .| .| .| o o'}, 'expected_output': [4, 2, 1, 2, 2, 1, 1, 1, 1, 4, 4]}
+                ,
+                '''
+def parse_music(music_string: str) -> List[int]:
+    """ Input to this function is a string representing musical notes in a special ASCII format.
+    Your task is to parse this string and return list of integers corresponding to how many beats does each
+    not last.
+
+    Here is a legend:
+    'o' - whole note, lasts four beats
+    'o|' - half note, lasts two beats
+    '.|' - quater note, lasts one beat
+
+    >>> parse_music('o o| .| o| o| .| .| .| .| o o')
+    [4, 2, 1, 2, 2, 1, 1, 1, 1, 4, 4]
+    """
+    note_map = {'o': 3, 'o|': 2, '.|': 1}
+    return [note_map[x] for x in music_string.split(' ') if x]
+    ''',
+            ],
+            [
+                {'input': {'numbers': [1.0, 2.0, 3.0], 'threshold': 0.5}, 'expected_output': False},
+                '''
+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
+    """
+    for idx, elem in enumerate(numbers):
+        for idx2, elem2 in enumerate(numbers):
+            if idx != idx2:
+                distance = elem - elem2
+                if distance < threshold:
+                    return True
+
+    return False
+    '''
+            ],
+            [
+                {'input': {'operations': [1, 2, -4, 5]}, 'expected_output': False},
+                '''
+def below_zero(operations: List[int]) -> bool:
+    """ You're given a list of deposit and withdrawal operations on a bank account that starts with
+    zero balance. Your task is to detect if at any point the balance of account fallls below zero, and
+    at that point function should return True. Otherwise it should return False.
+    >>> below_zero([1, 2, 3])
+    False
+    >>> below_zero([1, 2, -4, 5])
+    True
+    """
+    balance = 0
+
+    for op in operations:
+        balance += op
+        if balance == 0:
+            return True
+
+    return False
+    '''
+            ]
+        ],
+        inputs=[test_cases, code_input],
+    )
+
+
+app.launch()

testcase_utils.py

@@ -0,0 +1,222 @@
+from typing import List, Dict, Any
+from loguru import logger
+import ast
+import re
+import json
+from tqdm import tqdm
+
+
+def get_parameter_names(prompt: str, entry_point: str) -> List[str]:
+    """
+    Extract parameter names from the function signature in the prompt.
+    """
+    # logger.debug(f"Prompt: {prompt}")
+    # logger.debug(f"Entry point: {entry_point}")
+    tree = ast.parse(prompt)
+    for node in ast.walk(tree):
+        # logger.debug(f"Node name: {node.name if hasattr(node, 'name') else None}")
+        if isinstance(node, ast.FunctionDef) and node.name == entry_point:
+            # Return the parameter names from the function definition that matches the entry point
+            return [param.arg for param in node.args.args]
+    return []
+
+
+def parse_tests(test: str, parameter_names: List[str], entry_point: str) -> Dict[str, List[Dict[str, Any]]]:
+    """
+    Parse the test string into a structured format using AST.
+    """
+    # Remove the METADATA section
+    test = re.sub(r'METADATA = \{[^}]*\}', '', test)
+
+    # Parse the entire test string
+    tree = ast.parse(test)
+
+    test_cases = []
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Assert):
+            # Process each assert statement
+            test_case = process_assert(node, entry_point, parameter_names)
+            if test_case:
+                test_cases.append(test_case)
+
+    return {"test_cases": test_cases}
+
+
+def process_assert(node: ast.Assert, entry_point: str, parameter_names: List[str]) -> Dict[str, Any]:
+    """
+    Process a single assert statement and extract input and expected output.
+    """
+    if isinstance(node.test, ast.Compare) and isinstance(node.test.ops[0], ast.Eq):
+        left = node.test.left
+        right = node.test.comparators[0]
+
+        if isinstance(left, ast.Call) and isinstance(left.func, ast.Name) and left.func.id == "candidate":
+            input_dict = process_input(left.args, parameter_names)
+            # logger.debug(f"Input: {input_dict}")
+            # logger.debug(f"right: {right}")
+            # logger.debug(f"right type: {type(right)}")
+            # logger.debug(f"right value: {right.name if isinstance(right, ast.Name) else right.s if isinstance(right, ast.Str) else None}")
+
+            try:
+                # Attempt to evaluate using literal_eval
+                expected_output = ast.literal_eval(right)
+            except ValueError:
+                # Fallback to eval if literal_eval fails
+                # logger.warning("Falling back to eval due to failure in literal_eval")
+                expected_output = eval(compile(ast.Expression(right), filename="<ast>", mode="eval"))
+
+            return {"input": input_dict, "expected_output": expected_output}
+
+    return None
+
+
+def process_input(args: List[ast.expr], parameter_names: List[str]) -> Dict[str, Any]:
+    """
+    Process the input arguments and match them with parameter names.
+    """
+    input_dict = {}
+
+    for i, arg in enumerate(args):
+        try:
+            # Attempt to evaluate using literal_eval for simpler cases
+            evaluated_arg = ast.literal_eval(arg)
+        except ValueError:
+            # Fallback to eval if literal_eval fails
+            # logger.warning("Falling back to eval due to failure in literal_eval")
+            evaluated_arg = eval(compile(ast.Expression(arg), filename="<ast>", mode="eval"))
+
+        if i < len(parameter_names):
+            input_dict[parameter_names[i]] = evaluated_arg
+        else:
+            # Handle extra arguments if any
+            input_dict[f"arg_{i}"] = evaluated_arg
+
+    return input_dict
+
+
+def parse_all_problems(problems):
+    success_count = 0
+    unhandled_failures = 0
+    for problem in problems:
+        try:
+            problem = json.loads(problem)
+
+            # logger.info(f"Problem: {problem}")
+            # logger.debug(f"Test: {problem['test']}")
+
+            entry_point = problem["entry_point"]
+            parameter_names = get_parameter_names(problem["prompt"], entry_point)
+            # logger.info(f"Parameter names: {parameter_names}")
+
+            given_tests_raw = "\n".join(problem["given_tests"]).replace(entry_point, "candidate")
+            given_tests = parse_tests(given_tests_raw, parameter_names, entry_point)
+
+            # Parse the test cases using the parameter names
+            parsed_tests = parse_tests(problem["test"], parameter_names, entry_point)
+            # logger.info(f"Parsed tests: {parsed_tests}")
+            success_count += 1
+        except:
+            logger.exception(f"Error processing problem {problem['task_id']}")
+            if problem['is_solved'] == False:
+                unhandled_failures += 1
+            continue
+
+    logger.info(f"Success count: {success_count}")
+    logger.info(f"Total problems: {len(problems)}")
+    logger.info(f"Unhandled failures: {unhandled_failures}")
+
+
+def parse_specific_problem(problem):
+    try:
+        if isinstance(problem, str):
+            problem = json.loads(problem)
+
+        logger.info(f"Problem: {problem}")
+        logger.debug(f"Test: {problem['test']}")
+        logger.debug(f"Given Test: {problem['given_tests']}")
+
+        entry_point = problem["entry_point"]
+        parameter_names = get_parameter_names(problem["prompt"], entry_point)
+        logger.debug(f"Parameter names: {parameter_names}")
+
+        given_tests_raw = "\n".join(problem["given_tests"]).replace(entry_point, "candidate")
+        given_tests = parse_tests(given_tests_raw, parameter_names, entry_point)
+        logger.debug(f"Given tests: {given_tests}")
+
+        # Parse the test cases using the parameter names
+        all_tests = parse_tests(problem["test"], parameter_names, entry_point)
+        logger.debug(f"Parsed tests: {all_tests}")
+        return all_tests
+    except:
+        logger.exception(f"Error processing problem {problem['task_id']}")
+        return None
+
+#assert next_smallest([]) is None
+#assert decode_cyclic(encode_cyclic("abc")) == "abc"
+#assert round(find_zero([-6, 11, -6, 1]), 2) == 1.0
+#assert abs(candidate(1.33) - 0.33) < 1e-6
+
+def check_all_problems(problems):
+    problems_q = []
+    success_count = 0
+    fail_count = 0
+    for problem in tqdm(problems):
+        try:
+            problem = json.loads(problem)
+
+            logger.info(f"Problem: {problem}")
+            logger.debug(f"Test: {problem['test']}")
+            logger.debug(f"All Test: {problem['given_tests']}")
+
+            entry_point = problem["entry_point"]
+            parameter_names = get_parameter_names(problem["prompt"], entry_point)
+            logger.info(f"Parameter names: {parameter_names}")
+
+            # given_tests_len = len(problem["given_tests"])
+            # given_tests_raw = "\n".join(problem["given_tests"]).replace(entry_point, "candidate")
+            # given_tests = parse_tests(given_tests_raw, parameter_names, entry_point)
+            # parsed_given_tests_len = len(given_tests['test_cases'])
+            # assert given_tests_len == parsed_given_tests_len
+            # success_count += 1
+
+            #Parse the test cases using the parameter names
+            tests_len_candidate =  problem["test"].count('candidate')
+            parsed_tests = parse_tests(problem["test"], parameter_names, entry_point)
+            parsed_test_len = len(parsed_tests['test_cases'])
+            #assert parsed_test_len != 0 
+            assert tests_len_candidate - 1 == parsed_test_len
+            logger.info(f"Parsed tests: {parsed_tests}")
+            success_count += 1
+        except:
+            logger.exception(f"Error processing problem {problem['task_id']}")
+            if problem['is_solved'] == False:
+                fail_count += 1
+                problems_q.append(problem['task_id'])
+            continue
+    
+    with open('output_data/humaneval/seed/deepseek-coder-v2-lite-instruct/20240828-174550/dscoder_debugged_seeds_deepseek-coder-v2-lite-instruct_1_1_10.jsonl', "r") as f:
+        fixed = f.readlines()
+    for fix_problem in fixed:
+        fix_problem = json.loads(fix_problem)
+        if fix_problem['task_id'] in problems_q:
+            print(1)
+
+    logger.info(f"Success count: {success_count}")
+    logger.info(f"Total problems: {len(problems)}")
+    logger.info(f"Unhandled failures: {fail_count}")
+
+if __name__ == "__main__":
+    input_seeds = "input_data/humaneval/seed/deepseek-coder-v2-lite-instruct/seed.jsonl"
+
+    with open(input_seeds, "r") as f:
+        problems = f.readlines()
+
+    check_all_problems(problems)
+    #parse_all_problems(problems)
+
+    # parse the one with 'task_id': 'HumanEval/32'
+    # for problem in problems:
+    #     problem = json.loads(problem)
+    #     if problem['task_id'] == 'HumanEval/33':
+    #         parsed_tests = parse_specific_problem(problem)
+    #         break

timeout_utils.py

@@ -0,0 +1,67 @@
+import json
+import os
+from threading import Thread
+import resource
+import signal
+import json
+import os
+import traceback
+from loguru import logger
+
+
+def set_memory_limit(max_memory):
+    def limit_memory():
+        resource.setrlimit(resource.RLIMIT_AS, (max_memory, max_memory))
+    return limit_memory
+
+
+def timeout_handler(_, __):
+    raise TimeoutError()
+
+
+def to_jsonl(dict_data, file_path):
+    with open(file_path, 'a') as file:
+        json_line = json.dumps(dict_data)
+        file.write(json_line + os.linesep)
+
+
+class PropagatingThread(Thread):
+    def run(self):
+        self.exc = None
+        try:
+            if hasattr(self, '_Thread__target'):
+                # Thread uses name mangling prior to Python 3.
+                self.ret = self._Thread__target(*self._Thread__args, **self._Thread__kwargs)
+            else:
+                self.ret = self._target(*self._args, **self._kwargs)
+        except Exception as e:
+            self.exc = e
+
+    def join(self, timeout=None):
+        super(PropagatingThread, self).join(timeout)
+        if self.exc:
+            raise self.exc
+        if self.is_alive():
+            return None
+        return self.ret
+
+    def terminate(self):
+        self._stop()
+
+
+def function_with_timeout(func, args, timeout, max_memory=100 * 1024 * 1024):
+    result_container = []
+
+    def wrapper():
+        # set_memory_limit(max_memory)()
+        result_container.append(func(*args))
+
+    thread = PropagatingThread(target=wrapper)
+    thread.start()
+    thread.join(timeout)
+    if thread.is_alive():
+        logger.error(f"Timeout Error\n {args[0]} with timeout {timeout}")
+        thread.terminate()
+        raise TimeoutError()
+    else:
+        return result_container[0]

utils.py

@@ -0,0 +1,549 @@
+import ast
+from loguru import logger
+import sys
+import io
+import json
+import re
+import traceback
+import os
+from timeout_utils import function_with_timeout
+
+helpers = [
+    "import math",
+    "import re",
+    "import sys",
+    "import copy",
+    "import datetime",
+    "import itertools",
+    "import collections",
+    "import heapq",
+    "import statistics",
+    "import functools",
+    "import hashlib",
+    "import numpy",
+    "import numpy as np",
+    "import string",
+    "from typing import *",
+    "from collections import *",
+    "import heapq as hq",
+    "from itertools import *",
+    "from math import *",
+    "from statistics import *",
+    "from functools import *",
+    "from collections import *",
+    "from datetime import *",
+    "from copy import *",
+]
+
+STARTING_CODE = "\n".join(helpers)
+
+
+def create_dependency_graph(functions):
+    graph = {func_name: set() for func_name in functions}
+    for func_name, func_code in functions.items():
+        for other_func in functions:
+            if other_func in func_code and other_func != func_name:
+                graph[func_name].add(other_func)
+    return graph
+
+
+def topological_sort(graph):
+    visited = set()
+    stack = []
+
+    def dfs(node):
+        visited.add(node)
+        for neighbor in graph[node]:
+            if neighbor not in visited:
+                dfs(neighbor)
+        stack.append(node)
+
+    for node in graph:
+        if node not in visited:
+            dfs(node)
+
+    return stack
+
+
+def merge_changes_to_parents(func_name, dependency_graph, functions):
+    # Update the function in the functions dictionary
+    logger.info(f"Updating function {func_name} in the functions dictionary")
+
+    # For any function that calls the modified function, update its code
+    for parent, children in dependency_graph.items():
+        if func_name in children:
+            parent_code = functions[parent]
+            updated_parent_code = parent_code.replace(func_name, f"{func_name}")
+            functions[parent] = updated_parent_code
+            logger.info(f"Updated references to {func_name} in parent function {parent}")
+
+    # Regenerate the full code
+    full_code = "\n\n".join(functions.values())
+
+    logger.info(f"Merged changes from {func_name} to all relevant functions")
+    return full_code
+
+
+def extract_functions(code):
+    logger.info("Extracting functions from code")
+    tree = ast.parse(code)
+    functions = {}
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef):
+            func_code = ast.get_source_segment(code, node)
+            functions[node.name] = func_code
+    logger.info(f"Extracted {len(functions)} functions: {', '.join(functions.keys())}")
+    return functions
+
+
+def extract_code_blocks(response):
+    """Extract all code blocks from the response."""
+    return re.findall(r'```python\s*(.*?)\s*```', response, re.DOTALL)
+
+
+def extract_function(code_block, function_name):
+    """Extract a specific function from a code block."""
+    try:
+        tree = ast.parse(code_block)
+    except:
+        logger.error(f"Failed to parse code block for function: {function_name} from\n{code_block}")
+        return None
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef) and node.name == function_name:
+            return ast.get_source_segment(code_block, node)
+    return None
+
+
+def evaluate_given_tests(code, given_tests, max_memory=100 * 1024 * 1024):
+    test_code = f"{STARTING_CODE}\n\n{code}\n\n{given_tests}"
+    try:
+        function_with_timeout(exec, (test_code, globals()), timeout=10, max_memory=max_memory)
+        return True
+    except TimeoutError as e:
+        logger.error(f"Timeout Error: {str(e)}")
+    except MemoryError as e:
+        logger.error(f"Memory Error: {str(e)}")
+    except AssertionError as e:
+        logger.error(f"Assertion Error: {str(e)}")
+    except Exception as e:
+        logger.error(f'Error: {str(e)}')
+        logger.error(f'Traceback: {traceback.format_exc()}')
+    return False
+
+def evaluate_simple(code, entry_point, all_test, max_memory=100 * 1024 * 1024):
+    '''
+    directly concatenate the code and test code to evaluate on the private test cases
+    '''
+
+    test_code = f"{STARTING_CODE}\n\n{code}\n\n{all_test}\n\ncheck({entry_point})"
+    try:
+        function_with_timeout(exec, (test_code, globals()), timeout=10, max_memory=max_memory)
+        return True
+    except TimeoutError as e:
+        logger.error(f"Timeout Error: {str(e)}")
+    except MemoryError as e:
+        logger.error(f"Memory Error: {str(e)}")
+    except AssertionError as e:
+        logger.error(f"Assertion Error: {str(e)}")
+    except Exception as e:
+        logger.error(f'Error: {str(e)}')
+        logger.error(f'Traceback: {traceback.format_exc()}')
+    return False
+
+
+def evaluate(code, entry_point, testcase, return_trace=False):
+    logger.info(f"Evaluating {entry_point} with testcase: {testcase['input']}")
+
+    # Extract all functions from the code
+    try:
+        functions = extract_functions(code)
+    except:
+        logger.error(f"Failed to extract functions from code {code}")
+        # import pdb
+        # pdb.set_trace()
+    logger.info(f"Extracted functions: {', '.join(functions.keys())}")
+
+    # filter the functions that are called in the entry_point function
+    entry_point_function = functions[entry_point]
+    # entry_point_tree = ast.parse(entry_point_function)
+    # entry_point_calls = [node.func.id for node in ast.walk(entry_point_tree) if isinstance(node, ast.Call)]
+    # functions = {name: func for name, func in functions.items() if name in entry_point_calls}
+    # directly search for the string
+    functions = {name: func for name, func in functions.items() if name in entry_point_function}
+    logger.info(f"Filtered functions: {', '.join(functions.keys())}")
+
+    # Combine all functions into a single code block
+    full_code = "\n\n".join(functions.values())
+    # logger.info(f"Code being evaluated:\n{full_code}")
+
+    # Convert the input to a string representation that can be safely evaluated
+    input_repr = repr(testcase['input'])
+
+    if isinstance(testcase['input'], dict):
+        # Sometimes the input is a dictionary, which needs to be unpacked as keyword arguments
+        test_code = f'''{full_code}\n\nprint(repr({entry_point}(**{input_repr})))'''
+    else:
+        test_code = f'''{full_code}\n\nprint(repr({entry_point}({input_repr})))'''
+
+    # add the starting code to the test code
+    test_code = f"{STARTING_CODE}\n\n{test_code}"
+
+    old_stdout = sys.stdout
+    new_stdout = io.StringIO()
+    sys.stdout = new_stdout
+
+    try:
+        function_with_timeout(exec, (test_code, globals()), timeout=10)
+        output = new_stdout.getvalue().strip()
+        sys.stdout = old_stdout
+
+        # Convert both expected and actual output to the same type for comparison
+        expected_output = repr(testcase["expected_output"])
+
+        # Update actual_output before assertion
+        testcase['actual_output'] = ast.literal_eval(output)
+
+        assert output == expected_output, f"Expected {expected_output}, but got {output}"
+        logger.info(f'Test case passed: {testcase}')
+        logger.info(f'Expected: {expected_output}, Got: {output}')
+        return True, testcase
+    except TimeoutError:
+        logger.error(f'Test case failed: {testcase}')
+        logger.error(f"Timeout Error: {str(e)}")
+    except AssertionError as e:
+        logger.error(f'Test case failed: {testcase}')
+        logger.error(str(e))
+    except Exception as e:
+        logger.error(f'Test case failed: {testcase}')
+        logger.error(f'Error: {str(e)}')
+        logger.error(f'Traceback: {traceback.format_exc()}')
+        testcase['actual_output'] = str(e)
+        if return_trace:
+            testcase['traceback'] = traceback.format_exc()
+    finally:
+        sys.stdout = old_stdout
+
+    return False, testcase
+
+
+def extract_json_from_string(s):
+
+    # search for all the ```json blocks
+    matches = re.findall(r'```json\s*(.*?)\s*```', s, re.DOTALL)
+    if matches:
+        return matches[-1]
+    return None
+
+
+def parse_json_response(response):
+    json_str = extract_json_from_string(response)
+    if json_str:
+        try:
+            # Standard JSON corrections
+            json_str = json_str.strip().replace("True", "true")
+            json_str = json_str.replace("False", "false")
+            json_str = json_str.replace("'", '"')
+            json_str = json_str.replace("None", "null")
+
+            # Convert tuple notation to list notation
+            json_str = re.sub(r'\((-?\d+),\s*(-?\d+)\)', r'[\1, \2]', json_str)
+
+            logger.info(f"Extracted JSON string: {json_str}")
+            try:
+                return json.loads(json_str)
+            except:
+                # remove comments (for mistral model)
+                json_str = re.sub(r'#.*', '', json_str)
+                return json.loads(json_str)
+        except json.JSONDecodeError as e:
+            logger.error(f"Failed to parse extracted JSON: {json_str}")
+            logger.error(f"JSONDecodeError: {str(e)}")
+            # import pdb
+            # pdb.set_trace()
+    else:
+        logger.error("No JSON object found in the response")
+    return None
+
+
+def get_dependency_graph_str(graph, root=None, prefix="", is_last=True):
+    result = []
+
+    if root is None:
+        # Collect all roots if no specific root is given
+        roots = [node for node in graph if not any(node in children for children in graph.values())]
+        for i, root in enumerate(roots):
+            result.append(get_dependency_graph_str(graph, root, "", i == len(roots) - 1))
+        return "\n".join(result)
+
+    connector = "└── " if is_last else "├── "
+    result.append(prefix + connector + root)
+
+    if root in graph:
+        children = sorted(graph[root])
+        new_prefix = prefix + ("    " if is_last else "│   ")
+        for i, child in enumerate(children):
+            is_last_child = (i == len(children) - 1)
+            result.append(get_dependency_graph_str(graph, child, new_prefix, is_last_child))
+
+    return "\n".join(result)
+
+
+def extract_functions_from_code(node, parent=None):
+    """ Recursively extract functions and set parents. """
+    if isinstance(node, ast.Module):
+        for n in node.body:
+            extract_functions_from_code(n, parent=node)
+    elif isinstance(node, ast.FunctionDef):
+        node.parent = parent
+        if parent is not None and isinstance(parent, (ast.FunctionDef, ast.Module)):
+            parent.children.append(node)
+        for n in node.body:
+            extract_functions_from_code(n, parent=node)
+
+
+def split_nested_functions(code):
+    tree = ast.parse(code)
+    for node in ast.walk(tree):
+        node.children = []
+    extract_functions_from_code(tree)
+
+    flat_functions = []
+
+    def flatten_functions(node):
+        if isinstance(node, ast.FunctionDef):
+            flat_functions.append(node)
+            # Remove nested function definitions from the body
+            node.body = [n for n in node.body if not isinstance(n, ast.FunctionDef)]
+        for child in node.children:
+            flatten_functions(child)
+
+    flatten_functions(tree)
+
+    # Function to correct indentation for function docstrings
+    def correct_indentation(functions):
+        for func in functions:
+            # Get existing docstring if present
+            docstring = ast.get_docstring(func)
+            if docstring:
+                # Replace existing docstring node with corrected indentation
+                corrected_docstring = "\n".join([line if line.strip() != "" else "" for line in docstring.split("\n")])
+                func.body[0].value.s = corrected_docstring
+
+    correct_indentation(flat_functions)
+
+    return '\n\n'.join(ast.unparse(f).strip() for f in flat_functions)
+
+
+def remove_unused_functions(code, entry_point):
+
+    tree = ast.parse(code)
+
+    function_names = {node.name for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)}
+    function_calls = set()
+
+    class FunctionCallVisitor(ast.NodeVisitor):
+        def visit_Call(self, node):
+            if isinstance(node.func, ast.Name) and node.func.id in function_names:
+                function_calls.add(node.func.id)
+            self.generic_visit(node)
+
+    FunctionCallVisitor().visit(tree)
+
+    used_functions = set()
+
+    def mark_used(func_name):
+        if func_name not in used_functions:
+            used_functions.add(func_name)
+            for node in ast.walk(tree):
+                if isinstance(node, ast.FunctionDef) and node.name == func_name:
+                    FunctionCallVisitor().visit(node)
+                    for call in function_calls:
+                        mark_used(call)
+
+    mark_used(entry_point)
+
+    # only keep the functions that are used
+    tree.body = [node for node in tree.body if not isinstance(node, ast.FunctionDef) or node.name in used_functions]
+
+    all_unused_functions = function_names - used_functions
+
+    # convert back to code
+    return ast.unparse(tree), all_unused_functions
+
+
+def test_remove_unused_functions():
+    code = '''
+def rolling_max(numbers: List[int]) -> List[int]:
+    """From a given list of integers, generate a list of rolling maximum element found until given moment
+in the sequence.
+>>> rolling_max([1, 2, 3, 2, 3, 4, 2])
+[1, 2, 3, 3, 3, 4, 4]"""
+    (max_so_far, rolling_max_list) = initialize_max_and_list(numbers)
+    for num in numbers[1:]:
+        (max_so_far, rolling_max_list) = update_max_and_list(max_so_far, num, rolling_max_list)
+    return rolling_max_list
+
+def initialize_max_and_list(numbers: List[int]) -> Tuple[int, List[int]]:
+    max_so_far = numbers[0]
+    rolling_max_list = [max_so_far]
+    return (max_so_far, rolling_max_list)
+
+def update_max_and_list(max_so_far: int, num: int, rolling_max_list: List[int]) -> Tuple[int, List[int]]:
+    max_so_far = max(max_so_far, num)
+    rolling_max_list.append(max_so_far)
+    return (max_so_far, rolling_max_list)
+    
+def clean_data(data: List[str]) -> List[str]:
+    return [d.strip() for d in data]
+    '''.strip()
+
+    entry_point = "rolling_max"
+    logger.info(f"Original code:\n{code}")
+    output, unused_functions = remove_unused_functions(code, entry_point)
+    logger.info(f"Unused functions: {unused_functions}")
+    logger.info(f"Cleaned code:\n{output}")
+
+
+def test_split_nested_functions():
+    # The initial code provided by the user
+    code = '''
+def find_suffix_start(s: str) -> int:
+    for i in range(len(s)):
+        if is_palindrome(s[i:]):
+            return i
+    return 0
+
+def make_palindrome(string: str) -> str:
+    """This function takes a string and returns a palindrome by appending the reverse of the prefix of the string that makes it a palindrome."""
+    
+    
+    def is_palindrome(s: str) -> bool:
+        """
+        This function takes a string and returns True if it is a palindrome, False otherwise.
+        """
+    
+    
+        def compare(s: str) -> bool:
+            """
+            This function takes a string and returns True if it is a palindrome, False otherwise.
+            inner function
+            """
+
+            return s == s[::-1]
+    
+        return compare(s)
+
+    suffix_start = find_suffix_start(string)
+    return string + string[:suffix_start][::-1]
+    '''.strip()
+
+    # Splitting the nested functions and correcting the indentation
+    output = split_nested_functions(code)
+    print(output)
+
+
+def test_parse_json_response():
+
+    response = """
+**All Test Cases:**
+```json
+{
+    "test_cases": [
+        {"input": {"date": "03-11-2000"}, "expected_output": [11, 3, 2000]},
+        {"input": {"date": "15-01-2012"}, "expected_output": [15, 1, 2012]},
+        {"input": {"date": "04-0-2040"}, "expected_output": None},
+        {"input": {"date": "06-04-2020"}, "expected_output": [4, 6, 2020]},
+        {"input": {"date": "06/04/2020"}, "expected_output": None}
+    ]
+}
+```
+    """.strip()
+
+    parsed_json = parse_json_response(response)
+    print(parsed_json)
+
+
+def insert_docstring(code, docstring):
+
+    # surround the docstring with triple quotes
+    docstring = f'"""{docstring}"""'
+
+    lines = code.split('\n')
+    # Find the first non-empty line
+    first_line = next((i for i, line in enumerate(lines) if line.strip()), 0)
+
+    # Determine the indentation of the first line
+    indentation = len(lines[first_line]) - len(lines[first_line].lstrip())
+
+    # Find the 'def' line
+    def_line = next((i for i, line in enumerate(lines) if line.strip().startswith('def ')), first_line)
+
+    # Insert the docstring after the 'def' line, maintaining indentation
+    docstring_lines = [' ' * (indentation + 4) + line for line in docstring.split('\n')]
+    lines = lines[:def_line+1] + docstring_lines + lines[def_line+1:]
+
+    return '\n'.join(lines)
+
+
+def parse_transcoder_problem_content(problem):
+    # Extract the last group of content between [c++] and [python]
+    cpp_code = problem["prompt"].split("[c++]")[-1].split("[python]")[0].strip()
+    full_question = f'This function is translated into Python from the following C++ code: \n{cpp_code}\n'
+
+    try:
+        # Try to parse the existing solution
+        tree = ast.parse(problem["solution"])
+
+        # Create a new docstring node
+        docstring = ast.Expr(ast.Str(full_question))
+
+        # Find the first function definition in the AST
+        for node in tree.body:
+            if isinstance(node, ast.FunctionDef):
+                # Insert the docstring at the beginning of the function body
+                node.body.insert(0, docstring)
+                break
+        else:
+            # If no function definition is found, add the docstring at the end of the module
+            tree.body.append(docstring)
+
+        # Convert the modified AST back to source code
+        modified_solution = ast.unparse(tree)
+
+    except SyntaxError:
+        # If there's a syntax error, use the string-based method
+        logger.debug(f"Failed to parse solution for problem: {problem['task_id']}")
+        modified_solution = insert_docstring(problem["solution"], full_question)
+        logger.debug(f"Modified solution: {modified_solution}")
+
+    # Update the problem dictionary with the modified solution
+    problem["solution"] = modified_solution
+
+    return problem
+
+
+def test_parse_transcoder_problem_content():
+
+    input_seeds = "input_data/transcoder/seed/starcoder/seed.jsonl"
+    with open(input_seeds, "r") as f:
+        problems = [json.loads(line) for line in f]
+
+    for problem in problems:
+        try:
+            result = parse_transcoder_problem_content(problem)
+        except Exception as e:
+            logger.error(f"Failed to parse solution for problem: {problem['task_id']}")
+            logger.error(f"The solution is: \n{problem['solution']}")
+            logger.error(f"Error: {str(e)}")
+
+    logger.info("Successfully parsed all solutions")
+    # show an example
+    logger.info(f"Example result: {result['solution']}")
+
+
+if __name__ == "__main__":
+
+    # test_split_nested_functions()
+    # test_parse_json_response()
+    # test_remove_unused_functions()
+    test_parse_transcoder_problem_content()