import os

os.environ["TOKENIZERS_PARALLELISM"] = "false"
import sys
import argparse

import time

current_directory = os.getcwd()
module_path = os.path.abspath(os.path.join(current_directory, "src/"))
if module_path not in sys.path:
    sys.path.append(module_path)

from utils import *

from dotenv import load_dotenv

load_dotenv()
MODEL_PATH = os.getenv("MODEL_PATH")

from utils import load_model
from merge_funcs import *
from my_eval import (
    eval_df_ModAdd,
    eval_df_FineWeb,
    generate_samples_modadd_fineweb,
    concat_task_csv_files,
)
import nltk

def main(
    models_path,
    save_dir,
    model_size,
    block_list,
    evaluate_option,
    n_samples_modadd=10000,
    batch_size_modadd=400,
    n_sampled_fineweb=10000,
    batch_size_fineweb=32,
    save_dir_graft="GRAFT",
    save_dir_fine_tune="Fine_tune",
    verbose=True,
    vanilla_model_name=None,
    host_model_name=None,
    model_names=["Tuned Model", "Transformed Model", "Vanilla Model", "Final Model"],
):
    if vanilla_model_name is None:
        vanilla_model_name = f"EleutherAI/pythia-{model_size}M"
    if host_model_name is None:
        host_model_name = f"EleutherAI/pythia-{model_size}M-deduped"

    # Check if the directory already exists
    if not os.path.exists(os.path.join(models_path, save_dir)):
        os.makedirs(os.path.join(models_path, save_dir))

    tokenizer = AutoTokenizer.from_pretrained(vanilla_model_name)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"

    blocks_str = "_".join([str(x) for x in block_list])

    if verbose:
        print("Loading models...")

    for name in model_names:
        if verbose:
            print(f"Generating samples for {name}")

        model_dir = save_dir_fine_tune if "Tuned Model" in name else save_dir_graft
        model_path = os.path.join(models_path, model_dir)

        model = load_model(
            model_type=name,
            model_path=model_path,
            blocks_str=blocks_str,
            vanilla_model_name=vanilla_model_name,
            host_model_name=host_model_name,
        )
        model.generation_config.pad_token_id = tokenizer.pad_token_id

        sanitized_name = name.replace(" ", "_")
        footer = f"{blocks_str}_{sanitized_name}"

        output_df_modadd, output_df_fineweb = generate_samples_modadd_fineweb(
            models=[(model, sanitized_name)],
            tokenizer=tokenizer,
            footer=footer,
            model_path=models_path,
            save_dir=os.path.join(save_dir, sanitized_name),
            data_path=DATA_SAVE_PATH,
            n_samples_modadd=n_samples_modadd,
            batch_size_modadd=batch_size_modadd,
            max_samples_fineweb=n_sampled_fineweb,
            batch_size_fineweb=batch_size_fineweb,
            max_tokens_generated=30,
            mod=4,
        )

    ##########EVAL#########

    footer = f"{blocks_str}"

    if evaluate_option in ["modular_addition", "both"]:
        if verbose:
            print("Evaluating Modular Addition results...")
        all_model_generated_samples = concat_task_csv_files(
            os.path.join(models_path, save_dir),
            task="Modular_addition",
            blocks_str=blocks_str,
        )
        results_modadd = eval_df_ModAdd(
            all_model_generated_samples, return_mean_std=True
        )
        results_path = os.path.join(
            models_path, save_dir, f"Modular_addition_results_{footer}.csv"
        )
        results_modadd.to_csv(results_path)
        if verbose:
            print("Modular Addition evaluation completed.")

    if evaluate_option in ["fineweb", "both"]:
        if verbose:
            print("Evaluating FineWeb results...")
        all_model_generated_samples_fineweb = concat_task_csv_files(
            os.path.join(models_path, save_dir),
            task="FineWeb",
            blocks_str=blocks_str,
        )
        nltk.download("punkt")

        results_fineweb = eval_df_FineWeb(
            all_model_generated_samples_fineweb, return_mean_std=True
        )
        results_path_fineweb = os.path.join(
            models_path, save_dir, f"FineWeb_results_{footer}.csv"
        )
        results_fineweb.to_csv(results_path_fineweb)
        if verbose:
            print("FineWeb evaluation completed.")

if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description="Script to manage model merging and grafting."
    )
    parser.add_argument(
        "--models_path", type=str, default=MODEL_PATH, help="Model_path"
    )
    parser.add_argument(
        "--save_dir",
        type=str,
        default="samples_generated",
        help="Directory to save results generated by each model.",
    )
    parser.add_argument(
        "--save_dir_graft",
        type=str,
        default="GRAFT",
        help="Directory to save grafted models.",
    )
    parser.add_argument(
        "--save_dir_fine_tune",
        type=str,
        default="Fine_tune",
        help="Directory to save finetuned models.",
    )
    parser.add_argument(
        "--max_samples_modadd",
        type=int,
        default=1024,
        help="Maximum samples per grafting.",
    )
    parser.add_argument(
        "--max_samples_fineweb",
        type=int,
        default=50,
        help="Maximum samples per grafting.",
    )
    parser.add_argument(
        "--batch_size_modadd", type=int, default=30, help="Batch size for grafting."
    )
    parser.add_argument(
        "--batch_size_fineweb", type=int, default="70", help="Size of the Pythia model."
    )
    parser.add_argument(
        "--model_size", type=int, default="70", help="Size of the Pythia model."
    )
    parser.add_argument(
        "--block_list",
        type=lambda value: [int(x) for x in value.split(",")],
        default=[3],
        help="Number of layers",
    )
    parser.add_argument(
        "--evaluate",
        type=str,
        choices=["modular_addition", "fineweb", "both"],
        default="both",
        help="Specify which evaluation to perform: 'modular_addition', 'fineweb', or 'both'.",
    )
    parser.add_argument(
        "--host_model_name",
        type=str,
        default=f"EleutherAI/pythia-70M-deduped",
        help="host_model_name",
    )
    parser.add_argument(
        "--vanilla_model_name",
        type=str,
        default=f"EleutherAI/pythia-70M",
        help="vanilla_model_name",
    )
    args = parser.parse_args()

    main(
        models_path=args.models_path,
        save_dir=args.save_dir,
        save_dir_graft=args.save_dir_graft,
        save_dir_fine_tune=args.save_dir_fine_tune,
        n_samples_modadd=args.max_samples_modadd,
        batch_size_modadd=args.batch_size_modadd,
        n_sampled_fineweb=args.max_samples_fineweb,
        batch_size_fineweb=args.batch_size_fineweb,
        model_size=args.model_size,
        block_list=args.block_list,
        evaluate_option=args.evaluate,
        host_model_name=args.host_model_name,
        vanilla_model_name=args.vanilla_model_name,
    )

Instruction: add model_names as an argument

import inspect
import pytest
import argparse
import re
import sys
import types
from unittest.mock import patch, MagicMock, Mock

import inspect

# Create mock modules for any imported modules in the implementations
class MockAutoTokenizer:
    @staticmethod
    def from_pretrained(*args, **kwargs):
        mock = Mock()
        mock.pad_token = None
        mock.eos_token = "eos_token"
        mock.padding_side = None
        return mock

class MockUtils:
    @staticmethod
    def load_model(*args, **kwargs):
        mock_model = Mock()
        mock_model.generation_config = Mock()
        mock_model.generation_config.pad_token_id = None
        return mock_model

# Add mocks to sys.modules
@pytest.fixture(autouse=True)
def mock_dependencies(monkeypatch):
    # Mock utils module
    mock_utils = types.ModuleType("utils")
    mock_utils.load_model = MockUtils.load_model
    mock_utils.DATA_SAVE_PATH = "/mock/data/path"
    monkeypatch.setitem(sys.modules, "utils", mock_utils)

    # Mock merge_funcs module
    mock_merge_funcs = types.ModuleType("merge_funcs")
    monkeypatch.setitem(sys.modules, "merge_funcs", mock_merge_funcs)

    # Mock my_eval module
    mock_my_eval = types.ModuleType("my_eval")
    mock_my_eval.eval_df_ModAdd = Mock(return_value=Mock())
    mock_my_eval.eval_df_FineWeb = Mock(return_value=Mock())
    mock_my_eval.generate_samples_modadd_fineweb = Mock(return_value=(Mock(), Mock()))
    mock_my_eval.concat_task_csv_files = Mock(return_value=Mock())
    monkeypatch.setitem(sys.modules, "my_eval", mock_my_eval)

    # Mock nltk module
    mock_nltk = types.ModuleType("nltk")
    mock_nltk.download = Mock()
    monkeypatch.setitem(sys.modules, "nltk", mock_nltk)

    # Mock dotenv module
    mock_dotenv = types.ModuleType("dotenv")
    mock_dotenv.load_dotenv = Mock()
    monkeypatch.setitem(sys.modules, "dotenv", mock_dotenv)

    # Mock AutoTokenizer
    mock_transformers = types.ModuleType("transformers")
    mock_transformers.AutoTokenizer = MockAutoTokenizer
    monkeypatch.setitem(sys.modules, "transformers", mock_transformers)

    # Mock the os module
    mock_os = types.ModuleType("os")
    mock_os.environ = {"MODEL_PATH": "/mock/model/path", "TOKENIZERS_PARALLELISM": "false"}
    mock_os.path = MagicMock()
    mock_os.path.exists = Mock(return_value=True)
    mock_os.path.join = lambda *args: "/".join(args)
    mock_os.path.abspath = lambda path: path
    mock_os.getcwd = Mock(return_value="/mock/cwd")
    mock_os.makedirs = Mock()
    monkeypatch.setitem(sys.modules, "os", mock_os)

def test_model_names_parameter_exists(implementation):
    """Test that model_names parameter exists in the main function."""
    impl_name, module = implementation

    # Extract the source code
    source_code = inspect.getsource(module)

    # Check if main function is defined with model_names parameter
    main_func_match = re.search(r'def\s+main\s*\((.+?)\):', source_code, re.DOTALL)

    assert main_func_match, f"{impl_name} should define a main function"

    # Check if model_names is a parameter or has a default value
    param_list = main_func_match.group(1)
    assert "model_names" in param_list, f"{impl_name}'s main function should have a model_names parameter"

def test_model_names_default_value(implementation):
    """Test that the default value for model_names includes expected model types."""
    impl_name, module = implementation

    # Extract the main function definition
    source_code = inspect.getsource(module)
    main_func_match = re.search(r'def\s+main\s*\((.+?)\):', source_code, re.DOTALL)
    assert main_func_match, f"{impl_name} should define a main function"

    param_list = main_func_match.group(1)
    model_names_param = re.search(r'model_names\s*=\s*(\[.+?\])', param_list)

    if not model_names_param:
        pytest.fail(f"{impl_name}'s main function should have a default value for model_names")

    default_value = model_names_param.group(1)

    # Expected model names (at least these should be included)
    expected_models = ["Tuned Model", "Transformed Model", "Vanilla Model", "Final Model"]

    # Check if all expected models are in the default value
    for model in expected_models:
        assert model in default_value, f"{impl_name} should include '{model}' in default model names"

def test_command_line_arguments_setup(implementation):
    """Test that command line arguments are properly set up."""
    impl_name, module = implementation

    # Get the source code
    source_code = inspect.getsource(module)

    # Find the if __name__ == "__main__" block with a more robust pattern
    # This pattern will capture everything until the end of the file or the next function/class definition
    main_block_patterns = [
        r'if\s+__name__\s*==\s*[\'"]__main__[\'"]:(.+?)(?=\Z)',  # Match until end of file
        r'if\s+\*\*name\*\*\s*==\s*[\'"]__main__[\'"]:(.+?)(?=\Z)'  # Your pattern with **name**
    ]

    main_block = None
    for pattern in main_block_patterns:
        match = re.search(pattern, source_code, re.DOTALL)
        if match:
            main_block = match.group(1)
            break

    assert main_block is not None, f"{impl_name} should have a main block"

    # Print for debugging
    print(f"Main block found (first 100 chars): {main_block[:100]}...")
    print(f"Main block length: {len(main_block)}")

    # Check if ArgumentParser is used
    assert "ArgumentParser" in main_block, f"{impl_name} should use ArgumentParser in main block"

    # Check if main() is called anywhere in the main block
    assert "main(" in main_block.replace(" ", ""), f"{impl_name} should call main() in the main block"

def test_model_names_passed_to_main(implementation):
    """Test that model_names are passed to main function."""
    impl_name, module = implementation

    # Get the source code
    source_code = inspect.getsource(module)

    # Find the if __name__ == "__main__" block with improved pattern
    main_block_patterns = [
        r'if\s+__name__\s*==\s*[\'"]__main__[\'"]:(.+?)(?=\Z)',  # Match until end of file
        r'if\s+\*\*name\*\*\s*==\s*[\'"]__main__[\'"]:(.+?)(?=\Z)'  # Your pattern with **name**
    ]

    main_block = None
    for pattern in main_block_patterns:
        match = re.search(pattern, source_code, re.DOTALL)
        if match:
            main_block = match.group(1)
            break

    assert main_block is not None, f"{impl_name} should have a main block"

    # More robust way to find the main function call
    # This will find all the arguments passed to main()
    main_call_match = re.search(r'main\s*\((.*?)\s*\)', main_block, re.DOTALL)

    assert main_call_match, f"{impl_name} should call main() in the main block"

    main_args = main_call_match.group(1)

    # Different patterns for passing model_names
    patterns = [
        r'model_names\s*=\s*args\.model_names',    # model_names=args.model_names
        r'model_names\s*=\s*.*?model_names',       # model_names=some_var_with_model_names
        r'args\.model_names',                      # directly passing args.model_names
        r'model_names\s*=',                        # any assignment to model_names
        r'model_names\s*:',                        # model_names: value (alternative syntax)
    ]

    model_names_passed = any(re.search(pattern, main_args) for pattern in patterns)

    # If the regex patterns don't find it, do a simpler text search
    if not model_names_passed:
        model_names_passed = 'model_names' in main_args

    assert model_names_passed, f"{impl_name} should pass model_names to the main function. Found: {main_args}"

def test_model_names_used_in_loop(implementation):
    """Test that the model_names parameter is used in a loop in the main function."""
    impl_name, module = implementation

    # Extract the main function
    source_code = inspect.getsource(module)
    main_func_match = re.search(r'def\s+main\s*\(.+?\):(.*?)(?=\s*def|\s*if\s+__name__|\s*$|\Z)',
                              source_code, re.DOTALL)

    assert main_func_match, f"{impl_name} should define a main function"

    main_body = main_func_match.group(1)

    # Look for a loop over model_names
    has_loop = re.search(r'for\s+\w+\s+in\s+model_names', main_body) is not None

    assert has_loop, f"{impl_name}'s main function should iterate over model_names"

def test_model_name_used_in_function_calls(implementation):
    """Test that the model name from the loop is used in function calls."""
    impl_name, module = implementation

    # Extract the main function
    source_code = inspect.getsource(module)
    main_func_match = re.search(r'def\s+main\s*\(.+?\):(.*?)(?=\s*def|\s*if\s+__name__|\s*$|\Z)',
                              source_code, re.DOTALL)

    assert main_func_match, f"{impl_name} should define a main function"

    main_body = main_func_match.group(1)

    # Find the loop variable
    loop_var_match = re.search(r'for\s+(\w+)\s+in\s+model_names', main_body)

    assert loop_var_match, f"{impl_name}'s main function should have a clear loop over model_names"

    loop_var = loop_var_match.group(1)

    # Check if the loop variable is used meaningfully within the loop
    loop_start_pattern = f'for\\s+{loop_var}\\s+in\\s+model_names'
    loop_start_match = re.search(loop_start_pattern, main_body)
    if loop_start_match:
        # Find the portion of code after the loop start
        remaining_code = main_body[loop_start_match.end():]

        # Check if loop variable is used
        var_used = re.search(fr'{loop_var}\s*[=.,\(\)\[\]]', remaining_code) is not None
        assert var_used, f"{impl_name}'s main function should use the model name variable '{loop_var}' from the loop"

# def test_main_function_existence(implementation):
#     """Test that the main function exists with proper parameters."""
#     impl_name, module = implementation

#     # Check if main function exists
#     assert hasattr(module, 'main'), f"{impl_name} should have a main function"

#     # Check the signature of the main function
#     main_sig = inspect.signature(module.main)

#     # Required parameters
#     required_params = [
#         "models_path", "save_dir", "model_size", "block_list", "evaluate_option",
#     ]

#     for param in required_params:
#         assert param in main_sig.parameters, f"{impl_name}'s main function should have a '{param}' parameter"

def test_main_function_existence(implementation):
    """Test that the main function exists with proper parameters."""
    impl_name, module = implementation

    # First, check if the module was loaded properly or has errors
    if hasattr(module, '__error__'):
        pytest.skip(f"Module {impl_name} has errors: {module.__error__}")

    # Get the source code to manually check for main function definition
    source_code = inspect.getsource(module)

    # Check for main function definition using regex
    main_func_match = re.search(r'def\s+main\s*\(', source_code)
    assert main_func_match, f"{impl_name} should have a main function definition"

    # Now check if the module has the main attribute
    if not hasattr(module, 'main'):
        pytest.skip(f"{impl_name} has a main function definition but it couldn't be loaded")

    # If we get here, the main function exists, so check its parameters
    main_sig = inspect.signature(module.main)

    # Required parameters
    required_params = [
        "models_path", "save_dir", "model_size", "block_list", "evaluate_option",
    ]

    for param in required_params:
        assert param in main_sig.parameters, f"{impl_name}'s main function should have a '{param}' parameter"

@patch("argparse.ArgumentParser.parse_args")
def test_cli_args_handling(mock_args, implementation):
    """Test that CLI arguments are correctly handled and passed to main."""
    impl_name, module = implementation

    # Create a mock for parsed args
    args_mock = MagicMock()

    # Set required attributes
    args_mock.models_path = "test_path"
    args_mock.save_dir = "test_save_dir"
    args_mock.model_size = 70
    args_mock.block_list = [3]
    args_mock.evaluate = "both"
    args_mock.max_samples_modadd = 100
    args_mock.batch_size_modadd = 10
    args_mock.max_samples_fineweb = 50
    args_mock.batch_size_fineweb = 5
    args_mock.save_dir_graft = "GRAFT"
    args_mock.save_dir_fine_tune = "Fine_tune"
    args_mock.host_model_name = "host_model"
    args_mock.vanilla_model_name = "vanilla_model"
    args_mock.model_names = ["Model1", "Model2"]

    # Configure mock to return args
    mock_args.return_value = args_mock

    # Get the source code to check for main() call pattern
    source_code = inspect.getsource(module)

    # Find the if __name__ == "__main__" block with improved pattern
    main_block_patterns = [
        r'if\s+__name__\s*==\s*[\'"]__main__[\'"]:(.+?)(?=\Z)',  # Match until end of file
        r'if\s+\*\*name\*\*\s*==\s*[\'"]__main__[\'"]:(.+?)(?=\Z)'  # Your pattern with **name**
    ]

    main_block = None
    for pattern in main_block_patterns:
        match = re.search(pattern, source_code, re.DOTALL)
        if match:
            main_block = match.group(1)
            break

    assert main_block is not None, f"{impl_name} should have a main block"

    # Check ArgumentParser usage
    argparser_match = re.search(r'parser\s*=\s*argparse\.ArgumentParser', main_block)
    assert argparser_match, f"{impl_name} should create an ArgumentParser"

    # Find main call - use a more flexible regex pattern
    main_call_patterns = [
        r'main\s*\((.*?)\)',  # Simple pattern: main(...)
        r'main\s*\([^)]*\)',  # Handles multi-line arguments better
        r'main\s*\(([^)]*?)\)'  # Another attempt to capture args
    ]

    main_args = None
    for pattern in main_call_patterns:
        match = re.search(pattern, main_block, re.DOTALL)
        if match and len(match.groups()) > 0:
            main_args = match.group(1)
            break

    # If regex didn't work, try to find the call by other means
    if not main_args:
        # Find the position of 'main(' in the block
        main_pos = main_block.find('main(')
        if main_pos >= 0:
            # Extract from 'main(' to the matching ')'
            open_count = 1
            close_pos = main_pos + 5  # Start after 'main('
            while open_count > 0 and close_pos < len(main_block):
                if main_block[close_pos] == '(':
                    open_count += 1
                elif main_block[close_pos] == ')':
                    open_count -= 1
                close_pos += 1

            if open_count == 0:
                main_args = main_block[main_pos+5:close_pos-1]

    assert main_args is not None, f"{impl_name} should call main() in the main block"

    # Check essential parameters are passed
    essential_params = ["models_path", "save_dir", "model_size", "block_list", "model_names"]

    for param in essential_params:
        # Different patterns for passing parameters
        param_passed = (
            f"{param}=args.{param}" in main_args or
            f"{param}=" in main_args or
            f"args.{param}" in main_args
        )
        assert param_passed, f"{impl_name} should pass {param} to main()"

def test_arg_parser_for_model_names(implementation):
    """Test that ArgumentParser is configured to accept model_names."""
    import inspect
    import re

    impl_name, module = implementation

    # Get the source code
    source_code = inspect.getsource(module)

    # Find the if __name__ == "__main__" block with improved pattern
    main_block_patterns = [
        r'if\s+__name__\s*==\s*[\'"]__main__[\'"]:(.+?)(?=\Z)',  # Match until end of file
        r'if\s+\*\*name\*\*\s*==\s*[\'"]__main__[\'"]:(.+?)(?=\Z)'  # Your pattern with **name**
    ]

    main_block = None
    for pattern in main_block_patterns:
        match = re.search(pattern, source_code, re.DOTALL)
        if match:
            main_block = match.group(1)
            break

    assert main_block is not None, f"{impl_name} should have a main block"

    # Look for argument parser configuration for model names
    # Multiple patterns to catch different ways of defining the model_names argument
    model_args_patterns = [
        r'add_argument\(\s*[\'"]--model_names[\'"]',  # Standard format
        r'add_argument\(\s*"--model_names"',          # Double quotes
        r'add_argument\(\s*\'--model_names\'',        # Single quotes
        r'add_argument\([\'"]--model[-_]names[\'"]'   # Handle possible dash/underscore variation
    ]

    has_model_names_arg = False
    for pattern in model_args_patterns:
        if re.search(pattern, main_block):
            has_model_names_arg = True
            break

    assert has_model_names_arg, f"{impl_name} should have an ArgumentParser argument for model_names"

    # Check for lambda parsing of model_names (common pattern)
    lambda_pattern = r'type\s*=\s*lambda.*?split'
    uses_lambda_for_model_names = False

    if re.search(lambda_pattern, main_block):
        # Find the context around the lambda
        lambda_context = re.findall(r'.{0,50}' + lambda_pattern + r'.{0,50}', main_block)
        # Check if any lambda is used in model_names context
        for context in lambda_context:
            if 'model_names' in context.lower():
                uses_lambda_for_model_names = True
                break

    assert uses_lambda_for_model_names, f"{impl_name} should use a lambda function to parse model_names from a string"

def test_arg_parser_for_model_names(implementation):
    """Test that ArgumentParser is configured to accept model_names."""
    import inspect
    import re

    impl_name, module = implementation

    # Get the source code
    source_code = inspect.getsource(module)

    # Find the if __name__ == "__main__" block with improved pattern
    main_block_patterns = [
        r'if\s+__name__\s*==\s*[\'"]__main__[\'"]:(.+?)(?=\Z)',  # Match until end of file
        r'if\s+\*\*name\*\*\s*==\s*[\'"]__main__[\'"]:(.+?)(?=\Z)'  # Your pattern with **name**
    ]

    main_block = None
    for pattern in main_block_patterns:
        match = re.search(pattern, source_code, re.DOTALL)
        if match:
            main_block = match.group(1)
            break

    assert main_block is not None, f"{impl_name} should have a main block"

    # Look for argument parser configuration for model_names
    model_args_patterns = [
        r'add_argument\(\s*[\'"]--model_names[\'"]',  # Standard format
        r'add_argument\(\s*"--model_names"',          # Double quotes
        r'add_argument\(\s*\'--model_names\'',        # Single quotes
        r'add_argument\([\'"]--model[-_]names[\'"]'   # Handle possible dash/underscore variation
    ]

    has_model_names_arg = False
    for pattern in model_args_patterns:
        if re.search(pattern, main_block):
            has_model_names_arg = True
            break

    assert has_model_names_arg, f"{impl_name} should have an ArgumentParser argument for model_names"

    # Improved check for lambda in model_names argument
    # This approach looks for lambda within a reasonable proximity to --model_names
    model_names_arg_match = None
    for pattern in model_args_patterns:
        match = re.search(pattern, main_block)
        if match:
            # Get the position of the match
            pos = match.start()
            # Look for the end of this argument definition (next add_argument or end of main block)
            next_arg = re.search(r'add_argument', main_block[pos+10:])
            end_pos = next_arg.start() + pos + 10 if next_arg else len(main_block)
            # Extract the full argument definition
            model_names_arg_match = main_block[pos:end_pos]
            break

    # Check for lambda in the model_names argument definition
    assert model_names_arg_match and "lambda" in model_names_arg_match, \
        f"{impl_name} should use a lambda function to parse model_names from a string"

pytest
pytest-mock
nltk
python-dotenv
transformers
torch

import pytest
import os
import sys
import json
from typing import Dict, List, Optional, Any

# Import from local test_utils.py in the same directory
from test_utils import TestUtils, TestResultsManager

# Load all implementations in the current sandbox
implementations = TestUtils.load_all_implementations()
test_results = TestResultsManager()

@pytest.fixture(scope="session")
def sandbox_dir():
    """Fixture to provide the sandbox directory path."""
    return os.path.dirname(os.path.abspath(__file__))

@pytest.fixture(scope="session")
def sandbox_name():
    """Fixture to provide the sandbox name."""
    return os.path.basename(os.path.dirname(os.path.abspath(__file__)))

@pytest.fixture(scope="session")
def all_implementations():
    """Fixture to provide all implementations as a dictionary."""
    return implementations

@pytest.fixture(params=list(implementations.items()))
def implementation(request):
    """Fixture to provide each implementation to tests one at a time."""
    return request.param

@pytest.fixture(scope="session")
def results_manager():
    """Fixture to provide access to the test results manager."""
    return test_results

# Hook for collecting test results
@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    """Pytest hook to collect test results."""
    # Execute all other hooks to obtain the report object
    outcome = yield
    rep = outcome.get_result()

    # We're only interested in the call outcome
    if rep.when == "call":
        if hasattr(item, "callspec") and "implementation" in item.callspec.params:
            # Get implementation name and module
            impl_name, _ = item.callspec.params["implementation"]

            # Get test name
            test_name = item.nodeid.split("::")[-1]

            # Record result
            if rep.passed:
                test_results.record_result(impl_name, test_name, True)
            elif rep.failed:
                error_msg = str(rep.longrepr) if rep.longrepr else "Test failed"
                test_results.record_result(impl_name, test_name, False, error_msg)
            elif rep.skipped:
                skip_reason = rep.longrepr[2] if rep.longrepr else "Test skipped"
                test_results.record_skip(impl_name, test_name, skip_reason)

# Hook to save results at the end of testing
@pytest.hookimpl(trylast=True)
def pytest_sessionfinish(session, exitstatus):
    """Save test results at the end of the test session."""
    test_results.save_results()

import os
import sys
import glob
import re
import importlib.util
import traceback
import types
from typing import Dict, List, Optional, Any, Tuple

class TestUtils:
    @staticmethod
    def discover_implementation_files(directory: str = None) -> List[str]:
        """Find all implementation files in the current sandbox directory."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        patterns = [
            r'modified_code\d+\.py',
            r'new_code\d+\.py',
            r'original_code\.py',
            r'implementation\d*\.py'
        ]

        pattern = re.compile('|'.join(f'({p})' for p in patterns))
        implementations = []

        for file_path in glob.glob(os.path.join(directory, '*.py')):
            if pattern.search(os.path.basename(file_path)):
                implementations.append(file_path)

        # Sort files numerically
        def sort_key(path):
            filename = os.path.basename(path)
            match = re.search(r'(\d+)', filename)
            return int(match.group(1)) if match else 0

        return sorted(implementations, key=sort_key)

    @staticmethod
    def create_mock_module(file_path: str, module_name: str, error_info: str) -> types.ModuleType:
        """Create a mock module that contains error information but can still be tested."""
        # Create a new module object
        mock_module = types.ModuleType(module_name)

        # Add basic attributes
        mock_module.__file__ = file_path
        mock_module.__name__ = module_name
        mock_module.__display_name__ = module_name
        mock_module.__error__ = error_info

        # Add a dummy function that can be detected by test functions
        def dummy_function(*args, **kwargs):
            return f"Error in module: {error_info}"

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

    @staticmethod
    def load_module(file_path: str, module_name: Optional[str] = None) -> Any:
        """
        Safely load a module from a file path with proper error handling.
        If the module has errors, return a mock module that can still be tested.
        """
        if module_name is None:
            module_name = os.path.basename(file_path).replace('.py', '')

        # Create a unique module name to avoid conflicts
        sandbox_id = os.path.basename(os.path.dirname(file_path))
        unique_module_name = f"{sandbox_id}_{module_name}"

        try:
            # First, try to read the file to check for syntax errors
            with open(file_path, 'r') as f:
                source_code = f.read()

            # Check for syntax errors by compiling the code
            try:
                compiled = compile(source_code, file_path, 'exec')
            except SyntaxError as e:
                error_msg = f"Syntax error: {str(e)}"
                print(f"Syntax error in {file_path}: {e}")
                print(f"  Line {e.lineno}, column {e.offset}: {e.text}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module spec
            spec = importlib.util.spec_from_file_location(unique_module_name, file_path)
            if spec is None or spec.loader is None:
                error_msg = f"Could not create spec for {file_path}"
                print(f"Error: {error_msg}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module object
            module = importlib.util.module_from_spec(spec)
            sys.modules[unique_module_name] = module

            # Special handling for execution errors
            try:
                # Execute the module code in a safe way
                spec.loader.exec_module(module)
                # Store the original name for reference
                module.__display_name__ = module_name
                return module
            except Exception as e:
                error_msg = f"Runtime error: {str(e)}"
                traceback_str = traceback.format_exc()
                print(f"Error executing module {file_path}: {e}")
                print(traceback_str)

                # Create a partial module that contains what we loaded before the error
                mock_module = TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

                # Copy any attributes that might have been defined before the error
                for attr_name in dir(module):
                    if not attr_name.startswith('__'):
                        try:
                            setattr(mock_module, attr_name, getattr(module, attr_name))
                        except Exception:
                            pass  # Skip attributes that can't be copied

                return mock_module

        except FileNotFoundError as e:
            error_msg = f"File not found: {str(e)}"
            print(f"Error: {error_msg}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)
        except Exception as e:
            error_msg = f"Unexpected error: {str(e)}"
            print(f"Error loading module {file_path}: {e}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

    @classmethod
    def load_all_implementations(cls, directory: str = None) -> Dict[str, Any]:
        """Load all implementation files in the directory, including those with errors."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        implementations = {}

        implementation_files = cls.discover_implementation_files(directory)
        if not implementation_files:
            print("WARNING: No implementation files found. Check your file naming patterns.")

        for file_path in implementation_files:
            module_name = os.path.basename(file_path).replace('.py', '')
            module = cls.load_module(file_path, module_name)

            # Always add the module, even if it has errors
            implementations[module_name] = module

            if hasattr(module, '__error__'):
                print(f"Loaded with errors: {module_name} - {module.__error__}")
            else:
                print(f"Successfully loaded: {module_name}")

        return implementations

class TestResultsManager:
    def __init__(self):
        self.results = {}
        self.sandbox_name = os.path.basename(os.path.dirname(os.path.abspath(__file__)))

    def record_result(self, impl_name: str, test_name: str, passed: bool,
                     error_msg: Optional[str] = None) -> None:
        """Record a test result for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        if passed:
            self.results[impl_name]["passed"] += 1
        else:
            self.results[impl_name]["failed"] += 1
            if error_msg:
                self.results[impl_name]["errors"].append({
                    "test": test_name,
                    "error": error_msg
                })

    def record_skip(self, impl_name: str, test_name: str,
                   reason: Optional[str] = None) -> None:
        """Record a skipped test for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        self.results[impl_name]["skipped"] += 1
        if reason:
            self.results[impl_name]["errors"].append({
                "test": test_name,
                "error": f"SKIPPED: {reason}"
            })

    def get_winner(self) -> Tuple[Optional[int], Dict]:
        """Determine the winner based on test results."""
        winner = None
        max_passed = -1

        for impl_name, results in self.results.items():
            if impl_name == "original_code":
                continue  # Skip original code when determining winner

            if results["passed"] > max_passed:
                max_passed = results["passed"]
                winner = impl_name
            # Break ties by looking at failure count
            elif results["passed"] == max_passed and winner is not None:
                if results["failed"] < self.results[winner]["failed"]:
                    winner = impl_name

        # Convert winner to numeric index if possible
        winner_index = -1
        if winner and re.match(r'modified_code\d+', winner):
            try:
                winner_index = int(re.search(r'(\d+)', winner).group(1))
            except (AttributeError, ValueError):
                pass

        return winner_index, self.results

    def save_results(self, filename: str = "test_results.json") -> None:
        """Save test results to a JSON file."""
        import json

        winner_index, results = self.get_winner()

        # Check if all tests were skipped
        all_skipped = all(
            stats["skipped"] == stats["passed"] + stats["failed"] + stats["skipped"]
            for impl_name, stats in results.items()
            if impl_name != "original_code"
        )

        output = {
            "winner": winner_index,
            "all_skipped": all_skipped,
            "results": {
                name: {
                    "passed": stats["passed"],
                    "failed": stats["failed"],
                    "skipped": stats["skipped"],
                    "total": stats["passed"] + stats["failed"] + stats["skipped"]
                }
                for name, stats in results.items()
                if not name.startswith("_")  # Skip internal items
            }
        }

        with open(filename, "w") as f:
            json.dump(output, f, indent=2)

        print(f"Test results saved to {filename}")

        return output


from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, NoSuchElementException
import time
import sys
import win32gui
import win32con

def ocultar_janela_chrome():
    """Oculta a janela do Chrome usando win32gui"""
    def callback(hwnd, windows):
        if "chrome" in win32gui.GetWindowText(hwnd).lower():
            win32gui.ShowWindow(hwnd, win32con.SW_HIDE)
        return True

    win32gui.EnumWindows(callback, None)

no_number_button_path = '//*[@id="app"]/div/span[2]/div/span/div/div/div/div/div/div[2]/div/button'
# carregou //*[@id="app"]/div/div[2]/div[2]

# options.add_argument('--headless')
driver = webdriver.Chrome()
driver.get('https://web.whatsapp.com/')
# input('aa')

timer = 0
print('Aguardando o carregamento das conversas...')
while True:
    if timer > 180:
        sys.exit()

try:
        element = driver.find_element(
            (By.XPATH, '//*[@id="app"]/div/div[2]/div[2]')
        )
        if element:
            print(element.text)

        if element == 'Carregando suas conversas':
            break

    except Exception:
        pass

    timer = timer + 1
    time.sleep(1)

# input('Pressione Enter após fazer o login com QR code...')
print('LOGADO!')
time.sleep(5)  # Espera carregar a sessão

def verificar_numero_whatsapp(numero):
    try:
        url = f"https://web.whatsapp.com/send/?phone={numero}"
        driver.get(url)

        # Define um tempo máximo de espera
        wait = WebDriverWait(driver, 20)

        try:
            # Espera pelo elemento de chat ou mensagem de erro
            _ = wait.until(
                EC.presence_of_element_located(
                    (By.XPATH, '//*[@id="main"]/footer/div[1]/div/span/div/div[2]/div[1]')
                )
            )
            return True

        except TimeoutException:
            # Verifica se existe mensagem de erro
            try:
                _ = driver.find_element(By.XPATH, no_number_button_path)
                return False
            except NoSuchElementException:
                return False

    except Exception as ex:
        print(f'Erro inesperado: {ex}')
        return False

def verificar_lista_numeros(numeros):
    resultados = {}
    for numero in numeros:
        resultado = verificar_numero_whatsapp(numero)
        print(f'RESULTADO: {resultado}')
        resultados[numero] = resultado
        time.sleep(2)
    return resultados

# Exemplo de uso
if __name__ == "__main__":
    numeros_teste = [
        "1111111111111111"
    ]

    try:
        resultados = verificar_lista_numeros(numeros_teste)

        for numero, existe in resultados.items():
            status = "está" if existe else "não está"
            print(f"O número {numero} {status} registrado no WhatsApp")
    finally:
        driver.quit()  # Garante que o driver seja fechado ao finalizar

Instruction: whats wrong?

import pytest
from unittest.mock import patch, MagicMock, Mock
import inspect
import sys
import io
import re
import os
import importlib.util
from selenium.common.exceptions import NoSuchElementException

@pytest.fixture
def mock_driver():
    """Create a mock of webdriver.Chrome to avoid real browser interactions."""
    mock = MagicMock()
    mock.find_element.return_value = MagicMock()
    mock.get.return_value = None
    return mock

@pytest.fixture
def load_original_code():
    """Load the original code to compare with implementations."""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    original_path = os.path.join(script_dir, 'original_code.py')

    # If the original code file isn't in the test directory, use a hardcoded path
    if not os.path.exists(original_path):
        original_path = "/Users/waynechi/dev/copilot-arena-eval/experiments/sample_150/sandbox_748/original_code.py"

    spec = importlib.util.spec_from_file_location("original_module", original_path)
    original_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(original_module)
    return original_module

def test_error_in_find_element_syntax(implementation):
    """Test that the syntax error in find_element method is fixed."""
    impl_name, module = implementation

    # Get the source code to analyze the find_element call syntax
    source_code = inspect.getsource(module)

    # Check if find_element is called properly (not with a tuple as argument)
    improper_find_element = re.search(r'find_element\s*\(\s*\(', source_code)
    assert not improper_find_element, f"Implementation {impl_name} has improper find_element syntax with a tuple"

    # Check that find_element uses By.XPATH correctly
    proper_find_element = re.search(r'find_element\s*\(\s*By\.XPATH', source_code)
    assert proper_find_element, f"Implementation {impl_name} doesn't use By.XPATH correctly with find_element"

@patch('selenium.webdriver.Chrome')
def test_element_text_check(mock_chrome, implementation):
    """Test that element.text comparison is correctly implemented."""
    impl_name, module = implementation

    source_code = inspect.getsource(module)

    # Check that the code correctly accesses the text property of element
    element_text_access = re.search(r'element\.text', source_code)
    assert element_text_access, f"Implementation {impl_name} doesn't properly access the text property of the element"

    # Check for proper condition checking with element.text
    proper_text_check = (
        re.search(r'if\s+element\.text\s*==\s*[\'"]Carregando suas conversas[\'"]', source_code) or
        re.search(r'if\s+[\'"]Carregando suas conversas[\'"].*in\s+element\.text', source_code)
    )
    assert proper_text_check, f"Implementation {impl_name} doesn't properly check for 'Carregando suas conversas' text"

@patch('selenium.webdriver.Chrome')
def test_exception_handling(mock_chrome, implementation):
    """Test that exception handling is properly implemented."""
    impl_name, module = implementation

    source_code = inspect.getsource(module)

    # Check for specific exception handling
    specific_exception = re.search(r'except\s+NoSuchElementException', source_code)
    uses_specific_exception = bool(specific_exception)

    # If not using specific exception, check for general exception handling
    general_exception = re.search(r'except\s+Exception', source_code)
    has_exception_handling = uses_specific_exception or bool(general_exception)

    assert has_exception_handling, f"Implementation {impl_name} doesn't properly handle exceptions"

# @patch('selenium.webdriver.Chrome')
# def test_loading_conversations_loop(mock_chrome, implementation, monkeypatch):
#     """Test the loop that checks for loading conversations."""
#     impl_name, module = implementation

#     # Setup mocks
#     mock_element = MagicMock()
#     mock_element.text = "Carregando suas conversas"

#     # Mock find_element to return our mock element on the second call
#     find_element_calls = 0
#     def mock_find_element(*args, **kwargs):
#         nonlocal find_element_calls
#         find_element_calls += 1
#         if find_element_calls == 1:
#             raise NoSuchElementException()
#         return mock_element

#     mock_driver = MagicMock()
#     mock_driver.find_element = mock_find_element
#     mock_driver.get.return_value = None

#     # Patch time.sleep and sys.exit
#     monkeypatch.setattr('time.sleep', lambda x: None)
#     monkeypatch.setattr('sys.exit', lambda: None)

#     # Capture print outputs
#     captured_output = io.StringIO()
#     monkeypatch.setattr('sys.stdout', captured_output)

#     # Extract the while loop from the source code
#     source_code = inspect.getsource(module)
#     while_loop_pattern = re.compile(r'while True:.*?time\.sleep\(1\)', re.DOTALL)
#     while_loop_match = while_loop_pattern.search(source_code)

#     if while_loop_match:
#         loop_code = while_loop_match.group()
#         # Execute loop code in a controlled environment
#         try:
#             # Setup needed variables
#             timer = 0
#             driver = mock_driver

#             # Intercept the break command by raising a custom exception
#             class LoopBreak(Exception):
#                 pass

#             modified_loop = loop_code.replace('break', 'raise LoopBreak()')

#             try:
#                 exec(modified_loop,
#                      {'driver': mock_driver, 'timer': timer, 'time': MagicMock(),
#                       'sys': MagicMock(), 'By': MagicMock(), 'NoSuchElementException': NoSuchElementException,
#                       'LoopBreak': LoopBreak})
#             except LoopBreak:
#                 # Successfully broke out of the loop
#                 pass

#             # Check that the element's text was printed
#             output = captured_output.getvalue()
#             assert "Carregando suas conversas" in output, f"Implementation {impl_name} doesn't print element text"
#         except Exception as e:
#             pytest.fail(f"Failed to execute while loop code: {e}")

@patch('selenium.webdriver.Chrome')
def test_verificar_numero_whatsapp(mock_chrome, implementation):
    """Test that verificar_numero_whatsapp function works correctly."""
    impl_name, module = implementation

    # Check if the function exists
    assert hasattr(module, 'verificar_numero_whatsapp'), f"Implementation {impl_name} doesn't have verificar_numero_whatsapp function"

    # Get the source code
    func_source = inspect.getsource(module.verificar_numero_whatsapp)

    # Check for proper URL formatting
    url_format = re.search(r'url\s*=\s*f[\'"]https://web\.whatsapp\.com/send/\?phone=\{numero\}[\'"]', func_source)
    assert url_format, f"Implementation {impl_name} doesn't properly format WhatsApp URL"

    # Check for proper WebDriverWait usage
    wait_usage = re.search(r'WebDriverWait\(driver,\s*\d+\)', func_source)
    assert wait_usage, f"Implementation {impl_name} doesn't properly use WebDriverWait"

    # Check for proper exception handling
    exception_handling = re.search(r'except\s+(TimeoutException|Exception)', func_source)
    assert exception_handling, f"Implementation {impl_name} doesn't properly handle exceptions in verificar_numero_whatsapp"

@patch('selenium.webdriver.Chrome')
def test_verificar_lista_numeros(mock_chrome, implementation):
    """Test that verificar_lista_numeros function works correctly."""
    impl_name, module = implementation

    # Check if the function exists
    assert hasattr(module, 'verificar_lista_numeros'), f"Implementation {impl_name} doesn't have verificar_lista_numeros function"

    # Define a mock for verificar_numero_whatsapp
    with patch.object(module, 'verificar_numero_whatsapp', return_value=True) as mock_verify:
        # Call the function with test numbers
        test_numbers = ["1234567890", "0987654321"]
        results = module.verificar_lista_numeros(test_numbers)

        # Check that verificar_numero_whatsapp was called for each number
        assert mock_verify.call_count == len(test_numbers), f"Implementation {impl_name} doesn't call verificar_numero_whatsapp for each number"

        # Check that results are returned as a dictionary
        assert isinstance(results, dict), f"Implementation {impl_name} doesn't return a dictionary from verificar_lista_numeros"

        # Check that all test numbers are in the results
        for num in test_numbers:
            assert num in results, f"Implementation {impl_name} doesn't include all numbers in results"
            assert results[num] is True, f"Implementation {impl_name} doesn't correctly process results"

def test_driver_management(implementation):
    """Test that driver is properly initialized and closed."""
    impl_name, module = implementation

    source_code = inspect.getsource(module)

    # Check for driver initialization
    driver_init = re.search(r'driver\s*=\s*webdriver\.Chrome\(', source_code)
    assert driver_init, f"Implementation {impl_name} doesn't properly initialize the Chrome driver"

    # Check for driver quit in finally block
    driver_quit = re.search(r'finally:.*driver\.quit\(\)', source_code, re.DOTALL)
    assert driver_quit, f"Implementation {impl_name} doesn't properly quit the driver in a finally block"

def test_error_fixes_comparison(implementation, load_original_code):
    """Compare implementations to the original code to verify the error was fixed."""
    impl_name, module = implementation
    original_module = load_original_code

    # Get the original source code
    original_source = inspect.getsource(original_module)
    impl_source = inspect.getsource(module)

    # Check that the find_element error is fixed
    original_element_find = re.search(r'element = driver\.find_element\(\s*\(By\.XPATH', original_source)
    impl_element_find = re.search(r'element = driver\.find_element\(\s*By\.XPATH', impl_source)

    assert original_element_find and impl_element_find, f"Implementation {impl_name} didn't correctly fix the find_element syntax"
    assert not re.search(r'find_element\s*\(\s*\(', impl_source), f"Implementation {impl_name} still has wrong find_element syntax"

import pytest
import os
import sys
import json
from typing import Dict, List, Optional, Any

# Import from local test_utils.py in the same directory
from test_utils import TestUtils, TestResultsManager

# Load all implementations in the current sandbox
implementations = TestUtils.load_all_implementations()
test_results = TestResultsManager()

@pytest.fixture(scope="session")
def sandbox_dir():
    """Fixture to provide the sandbox directory path."""
    return os.path.dirname(os.path.abspath(__file__))

@pytest.fixture(scope="session")
def sandbox_name():
    """Fixture to provide the sandbox name."""
    return os.path.basename(os.path.dirname(os.path.abspath(__file__)))

@pytest.fixture(scope="session")
def all_implementations():
    """Fixture to provide all implementations as a dictionary."""
    return implementations

@pytest.fixture(params=list(implementations.items()))
def implementation(request):
    """Fixture to provide each implementation to tests one at a time."""
    return request.param

@pytest.fixture(scope="session")
def results_manager():
    """Fixture to provide access to the test results manager."""
    return test_results

# Hook for collecting test results
@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    """Pytest hook to collect test results."""
    # Execute all other hooks to obtain the report object
    outcome = yield
    rep = outcome.get_result()

    # We're only interested in the call outcome
    if rep.when == "call":
        if hasattr(item, "callspec") and "implementation" in item.callspec.params:
            # Get implementation name and module
            impl_name, _ = item.callspec.params["implementation"]

            # Get test name
            test_name = item.nodeid.split("::")[-1]

            # Record result
            if rep.passed:
                test_results.record_result(impl_name, test_name, True)
            elif rep.failed:
                error_msg = str(rep.longrepr) if rep.longrepr else "Test failed"
                test_results.record_result(impl_name, test_name, False, error_msg)
            elif rep.skipped:
                skip_reason = rep.longrepr[2] if rep.longrepr else "Test skipped"
                test_results.record_skip(impl_name, test_name, skip_reason)

# Hook to save results at the end of testing
@pytest.hookimpl(trylast=True)
def pytest_sessionfinish(session, exitstatus):
    """Save test results at the end of the test session."""
    test_results.save_results()

import os
import sys
import glob
import re
import importlib.util
import traceback
import types
from typing import Dict, List, Optional, Any, Tuple

class TestUtils:
    @staticmethod
    def discover_implementation_files(directory: str = None) -> List[str]:
        """Find all implementation files in the current sandbox directory."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        patterns = [
            r'modified_code\d+\.py',
            r'new_code\d+\.py',
            r'original_code\.py',
            r'implementation\d*\.py'
        ]

        pattern = re.compile('|'.join(f'({p})' for p in patterns))
        implementations = []

        for file_path in glob.glob(os.path.join(directory, '*.py')):
            if pattern.search(os.path.basename(file_path)):
                implementations.append(file_path)

        # Sort files numerically
        def sort_key(path):
            filename = os.path.basename(path)
            match = re.search(r'(\d+)', filename)
            return int(match.group(1)) if match else 0

        return sorted(implementations, key=sort_key)

    @staticmethod
    def create_mock_module(file_path: str, module_name: str, error_info: str) -> types.ModuleType:
        """Create a mock module that contains error information but can still be tested."""
        # Create a new module object
        mock_module = types.ModuleType(module_name)

        # Add basic attributes
        mock_module.__file__ = file_path
        mock_module.__name__ = module_name
        mock_module.__display_name__ = module_name
        mock_module.__error__ = error_info

        # Add a dummy function that can be detected by test functions
        def dummy_function(*args, **kwargs):
            return f"Error in module: {error_info}"

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

    @staticmethod
    def load_module(file_path: str, module_name: Optional[str] = None) -> Any:
        """
        Safely load a module from a file path with proper error handling.
        If the module has errors, return a mock module that can still be tested.
        """
        if module_name is None:
            module_name = os.path.basename(file_path).replace('.py', '')

        # Create a unique module name to avoid conflicts
        sandbox_id = os.path.basename(os.path.dirname(file_path))
        unique_module_name = f"{sandbox_id}_{module_name}"

        try:
            # First, try to read the file to check for syntax errors
            with open(file_path, 'r') as f:
                source_code = f.read()

            # Check for syntax errors by compiling the code
            try:
                compiled = compile(source_code, file_path, 'exec')
            except SyntaxError as e:
                error_msg = f"Syntax error: {str(e)}"
                print(f"Syntax error in {file_path}: {e}")
                print(f"  Line {e.lineno}, column {e.offset}: {e.text}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module spec
            spec = importlib.util.spec_from_file_location(unique_module_name, file_path)
            if spec is None or spec.loader is None:
                error_msg = f"Could not create spec for {file_path}"
                print(f"Error: {error_msg}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module object
            module = importlib.util.module_from_spec(spec)
            sys.modules[unique_module_name] = module

            # Special handling for execution errors
            try:
                # Execute the module code in a safe way
                spec.loader.exec_module(module)
                # Store the original name for reference
                module.__display_name__ = module_name
                return module
            except Exception as e:
                error_msg = f"Runtime error: {str(e)}"
                traceback_str = traceback.format_exc()
                print(f"Error executing module {file_path}: {e}")
                print(traceback_str)

                # Create a partial module that contains what we loaded before the error
                mock_module = TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

                # Copy any attributes that might have been defined before the error
                for attr_name in dir(module):
                    if not attr_name.startswith('__'):
                        try:
                            setattr(mock_module, attr_name, getattr(module, attr_name))
                        except Exception:
                            pass  # Skip attributes that can't be copied

                return mock_module

        except FileNotFoundError as e:
            error_msg = f"File not found: {str(e)}"
            print(f"Error: {error_msg}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)
        except Exception as e:
            error_msg = f"Unexpected error: {str(e)}"
            print(f"Error loading module {file_path}: {e}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

    @classmethod
    def load_all_implementations(cls, directory: str = None) -> Dict[str, Any]:
        """Load all implementation files in the directory, including those with errors."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        implementations = {}

        implementation_files = cls.discover_implementation_files(directory)
        if not implementation_files:
            print("WARNING: No implementation files found. Check your file naming patterns.")

        for file_path in implementation_files:
            module_name = os.path.basename(file_path).replace('.py', '')
            module = cls.load_module(file_path, module_name)

            # Always add the module, even if it has errors
            implementations[module_name] = module

            if hasattr(module, '__error__'):
                print(f"Loaded with errors: {module_name} - {module.__error__}")
            else:
                print(f"Successfully loaded: {module_name}")

        return implementations

class TestResultsManager:
    def __init__(self):
        self.results = {}
        self.sandbox_name = os.path.basename(os.path.dirname(os.path.abspath(__file__)))

    def record_result(self, impl_name: str, test_name: str, passed: bool,
                     error_msg: Optional[str] = None) -> None:
        """Record a test result for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        if passed:
            self.results[impl_name]["passed"] += 1
        else:
            self.results[impl_name]["failed"] += 1
            if error_msg:
                self.results[impl_name]["errors"].append({
                    "test": test_name,
                    "error": error_msg
                })

    def record_skip(self, impl_name: str, test_name: str,
                   reason: Optional[str] = None) -> None:
        """Record a skipped test for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        self.results[impl_name]["skipped"] += 1
        if reason:
            self.results[impl_name]["errors"].append({
                "test": test_name,
                "error": f"SKIPPED: {reason}"
            })

    def get_winner(self) -> Tuple[Optional[int], Dict]:
        """Determine the winner based on test results."""
        winner = None
        max_passed = -1

        for impl_name, results in self.results.items():
            if impl_name == "original_code":
                continue  # Skip original code when determining winner

            if results["passed"] > max_passed:
                max_passed = results["passed"]
                winner = impl_name
            # Break ties by looking at failure count
            elif results["passed"] == max_passed and winner is not None:
                if results["failed"] < self.results[winner]["failed"]:
                    winner = impl_name

        # Convert winner to numeric index if possible
        winner_index = -1
        if winner and re.match(r'modified_code\d+', winner):
            try:
                winner_index = int(re.search(r'(\d+)', winner).group(1))
            except (AttributeError, ValueError):
                pass

        return winner_index, self.results

    def save_results(self, filename: str = "test_results.json") -> None:
        """Save test results to a JSON file."""
        import json

        winner_index, results = self.get_winner()

        # Check if all tests were skipped
        all_skipped = all(
            stats["skipped"] == stats["passed"] + stats["failed"] + stats["skipped"]
            for impl_name, stats in results.items()
            if impl_name != "original_code"
        )

        output = {
            "winner": winner_index,
            "all_skipped": all_skipped,
            "results": {
                name: {
                    "passed": stats["passed"],
                    "failed": stats["failed"],
                    "skipped": stats["skipped"],
                    "total": stats["passed"] + stats["failed"] + stats["skipped"]
                }
                for name, stats in results.items()
                if not name.startswith("_")  # Skip internal items
            }
        }

        with open(filename, "w") as f:
            json.dump(output, f, indent=2)

        print(f"Test results saved to {filename}")

        return output


import os
import shutil

from transformers import AutoModelForCausalLM
from peft import PeftModel

from dotenv import load_dotenv

import pickle
import torch
import json

load_dotenv()

DATA_SAVE_PATH = os.getenv("DATA_SAVE_PATH")
MODEL_PATH = os.getenv("MODEL_PATH")

def save_log_to_file(log_history, file_path, append_latest_only=False):
    """
    Saves the log history to a JSON file.
    If the file already exists, it appends to it.

    Parameters:
    - log_history: List of log entries (each entry is a dict).
    - file_path: Path to the file where logs will be saved.
    - append_latest_only: If True, only the latest log entry is appended.
    """
    # Initialize current_logs
    current_logs = []

    # If the file exists, load the current logs and append to them
    if os.path.exists(file_path):
        try:
            with open(file_path, "r") as f:
                content = f.read().strip()
                if content:
                    current_logs = json.loads(content)
                else:
                    current_logs = []
        except json.JSONDecodeError:
            print(f"Warning: {file_path} contains invalid JSON. Overwriting file.")
            current_logs = []
        except Exception as e:
            print(f"An error occurred while reading {file_path}: {e}")
            current_logs = []
    else:
        # File does not exist; current_logs remains an empty list
        pass

    # Decide whether to append the entire log history or just the latest entry
    if append_latest_only and log_history:
        # Append only the most recent epoch log
        current_logs.append(log_history[-1])
    else:
        # Append the entire log history
        current_logs.extend(log_history)

    # Save the updated log history
    try:
        with open(file_path, "w") as f:
            json.dump(current_logs, f, indent=4)
    except Exception as e:
        print(f"An error occurred while writing to {file_path}: {e}")

def clear_directory(directory):
    """
    Clears all files and subdirectories within a given directory.  Creates the directory if it doesn't exist.

    Args:
        directory (str): The path to the directory to clear.

    Raises:
        OSError: If any error occurs during file or directory removal.  Provides details about the failure.
    Example:
        clear_directory('/path/to/my/directory')
    """
    if not os.path.exists(directory):
        os.makedirs(directory)
        print(f"Directory '{directory}' created.")
        return
    for item in os.listdir(directory):
        item_path = os.path.join(directory, item)
        try:
            if os.path.isdir(item_path):
                shutil.rmtree(item_path)
                print(f"Removed directory: {item_path}")
            else:
                os.remove(item_path)
                print(f"Removed file: {item_path}")
        except OSError as e:
            print(f"Failed to delete '{item_path}'. Reason: {e}")

def merge_lora_model(
    model_name="pythia-31M",
    base_model_repo_name="EleutherAI/",
    model_load_path=MODEL_PATH,
    model_save_path=MODEL_PATH,
):

    my_model_path = os.path.join(model_load_path, model_name)
    param_count = model_name.lower().split("m")[0].split("-")[1]
    base_model = f"pythia-{param_count}M"

    base_model = AutoModelForCausalLM.from_pretrained(
        os.path.join(base_model_repo_name, base_model)
    )
    model = PeftModel.from_pretrained(base_model, my_model_path)
    merged_model = model.merge_and_unload()
    my_model_save_path = os.path.join(model_save_path, f"{model_name}_merged")
    merged_model.save_pretrained(my_model_save_path)

def remove_repetition(question, answer):
    if question in answer:
        return answer.replace(question, "").strip()
    return answer

def load_model(
    model_type,
    model_path=None,
    blocks_str=None,
    vanilla_model_name=None,
    host_model_name=None,
):
    """
    Loads different types of models based on the model_type parameter.

    Parameters:
    model_type (str): The type of model to load. One of 'Tuned Model', 'Vanilla Model',
                      'Transformed Model', 'Final Model', or 'Host Model'.
    model_path (str): The base path where models are stored.
    blocks_str (str): A string representing the layers or blocks used in model naming.
    vanilla_model_name (str): The name or path of the vanilla (base) model.
    host_model_name (str): The name or path of the host model.

    Returns:
    model: The loaded model object.

    Raises:
    ValueError: If an unknown model_type is provided or required parameters are missing.
    IOError: If loading the model fails.

    Example:
    model = load_model(
        model_type="Tuned Model",
        model_path="/path/to/models",
        blocks_str="1-5",
        vanilla_model_name="EleutherAI/pythia-31M"
    )
    """
    if model_type == "Tuned Model":
        model_name = vanilla_model_name.split("/")[-1]

        # save_path = os.path.join(model_path)
        # model_save_name = f"{model_name}_trained_{footer}"
        # save_path = os.path.join(save_path, model_save_name)

        tuned_model_name = f"{model_name}_trained_layers_{blocks_str}_merged"
        tuned_model = AutoModelForCausalLM.from_pretrained(
            os.path.join(model_path, f"{tuned_model_name}")
        )
        return tuned_model

    elif model_type == "Vanilla Model":
        vanilla_model = AutoModelForCausalLM.from_pretrained(vanilla_model_name)
        return vanilla_model

    elif model_type == "Transformed Model":
        name = host_model_name.split("/")[-1]
        save_path = os.path.join(model_path, f"{name}_preGRAFTED_{blocks_str}.pkl")
        with open(save_path, "rb") as f:
            transformed_model = pickle.load(f)
        return transformed_model

    elif model_type == "Final Model":
        name = host_model_name.split("/")[-1]
        model_save_name = f"{name}_GRAFTED_{blocks_str}.pkl"
        save_path = os.path.join(model_path, model_save_name)
        with open(save_path, "rb") as f:
            final_model = pickle.load(f)
        return final_model
    elif model_type == "Host Model":
        host_model = AutoModelForCausalLM.from_pretrained(host_model_name)
        return host_model

    else:
        raise ValueError(f"Unknown model type: {model_type}")

def load_batch_losses(file_path):
    """
    Loads batch loss data from a checkpoint file.

    Parameters:
    file_path (str): The path to the checkpoint file.

    Returns:
    list or None: The batch losses if available, None otherwise.

    Logs:
    An error message if loading fails.

    Example:
    batch_losses = load_batch_losses('/path/to/checkpoint.pt')
    """
    try:
        checkpoint = torch.load(file_path, map_location=torch.device("cpu"))
        batch_losses = checkpoint.get("batch_losses", None)
        if batch_losses is not None:
            logging.info(f"Batch losses loaded from {file_path}")
        else:
            logging.warning(f"No 'batch_losses' key found in checkpoint at {file_path}")
        return batch_losses
    except (FileNotFoundError, IOError, RuntimeError) as e:
        logging.error(f"Error loading checkpoint from {file_path}: {e}")
        return None

Instruction: add the option to delete the whole directory

import os
import shutil
import inspect
import tempfile
import pytest
from unittest.mock import patch, MagicMock

def test_clear_directory_function_signature(implementation):
    """Test that clear_directory function has the required parameter for deletion."""
    impl_name, module = implementation

    # Check if the function has a parameter for deleting the directory
    sig = inspect.signature(module.clear_directory)

    # Get parameter names
    param_names = list(sig.parameters.keys())

    # Check if there's at least one parameter (directory)
    assert len(param_names) >= 1, f"Implementation {impl_name} should have at least 1 parameter"

    # If there are at least 2 parameters, check the deletion parameter
    if len(param_names) >= 2:
        # The deletion parameter name might vary, but should be the second parameter
        deletion_param = param_names[1]

        # Check that the parameter has a default value of False
        assert sig.parameters[deletion_param].default is False, \
            f"Implementation {impl_name} should have deletion parameter default to False"

def test_clear_directory_without_deletion(implementation):
    """Test clear_directory functions correctly when not deleting the directory."""
    impl_name, module = implementation

    with tempfile.TemporaryDirectory() as temp_dir:
        # Create some test files and subdirectories
        test_file_path = os.path.join(temp_dir, "test_file.txt")
        test_subdir_path = os.path.join(temp_dir, "test_subdir")

        with open(test_file_path, "w") as f:
            f.write("test content")

        os.makedirs(test_subdir_path)

        # Mock print function to avoid output during tests
        with patch('builtins.print'):
            # Clear directory without deletion flag
            module.clear_directory(temp_dir)

        # Directory should still exist
        assert os.path.exists(temp_dir)

        # Files and subdirectories should be removed
        assert len(os.listdir(temp_dir)) == 0

def test_clear_directory_with_deletion(implementation):
    """Test clear_directory function correctly deletes the entire directory."""
    impl_name, module = implementation

    with tempfile.TemporaryDirectory() as parent_dir:
        # Create a directory inside the temporary directory
        test_dir = os.path.join(parent_dir, "test_dir")
        os.makedirs(test_dir)

        # Create a test file
        test_file_path = os.path.join(test_dir, "test_file.txt")
        with open(test_file_path, "w") as f:
            f.write("test content")

        # Mock print function to avoid output during tests
        with patch('builtins.print'):
            # Get the parameter name for deletion
            sig = inspect.signature(module.clear_directory)
            param_names = list(sig.parameters.keys())

            # Check if implementation has a deletion parameter
            if len(param_names) < 2:
                pytest.skip(f"Implementation {impl_name} does not support deletion parameter")

            deletion_param = param_names[1]

            # Call clear_directory with deletion parameter set to True
            kwargs = {deletion_param: True}
            module.clear_directory(test_dir, **kwargs)

        # Directory should be deleted
        assert not os.path.exists(test_dir)

def test_clear_directory_creates_directory_if_nonexistent(implementation):
    """Test clear_directory creates the directory if it doesn't exist."""
    impl_name, module = implementation

    with tempfile.TemporaryDirectory() as parent_dir:
        # Define a non-existent directory path
        nonexistent_dir = os.path.join(parent_dir, "nonexistent_dir")

        # Make sure it doesn't exist
        if os.path.exists(nonexistent_dir):
            shutil.rmtree(nonexistent_dir)

        # Mock print function to avoid output during tests
        with patch('builtins.print'):
            # Call clear_directory on non-existent directory
            module.clear_directory(nonexistent_dir)

        # Directory should be created
        assert os.path.exists(nonexistent_dir)
        assert os.path.isdir(nonexistent_dir)

def test_clear_directory_with_deletion_no_recreation(implementation):
    """Test that clear_directory doesn't recreate directory after deletion."""
    impl_name, module = implementation

    with tempfile.TemporaryDirectory() as parent_dir:
        # Create a directory inside the temporary directory
        test_dir = os.path.join(parent_dir, "test_dir")
        os.makedirs(test_dir)

        # Create a test file
        test_file_path = os.path.join(test_dir, "test_file.txt")
        with open(test_file_path, "w") as f:
            f.write("test content")

        # Mock print function to avoid output during tests
        with patch('builtins.print'):
            # Get the parameter name for deletion
            sig = inspect.signature(module.clear_directory)
            param_names = list(sig.parameters.keys())

            # Skip test if implementation doesn't have a deletion parameter
            if len(param_names) < 2:
                pytest.skip(f"Implementation {impl_name} does not support deletion parameter")

            deletion_param = param_names[1]

            # Call clear_directory with deletion parameter set to True
            kwargs = {deletion_param: True}
            module.clear_directory(test_dir, **kwargs)

        # Directory should be deleted and not recreated
        assert not os.path.exists(test_dir)

def test_clear_directory_handles_errors_gracefully(implementation):
    """Test that clear_directory handles errors gracefully."""
    impl_name, module = implementation

    with tempfile.TemporaryDirectory() as temp_dir:
        # Create a test file
        test_file_path = os.path.join(temp_dir, "test_file.txt")
        with open(test_file_path, "w") as f:
            f.write("test content")

        # Mock os.remove to raise an OSError
        def mock_remove_with_error(*args, **kwargs):
            raise OSError("Mock error")

        # Mock necessary functions to ensure errors are caught
        with patch('os.remove', side_effect=mock_remove_with_error), \
             patch('builtins.print') as mock_print, \
             patch('os.rmdir'), patch('shutil.rmtree'):

            try:
                # Call clear_directory
                module.clear_directory(temp_dir)
                # If we reach here, the function caught the error
                assert mock_print.called, "Function should print an error message"
            except OSError:
                # If OSError was raised, check if it was at least logged
                assert mock_print.called, "Function should print an error before raising"

def test_clear_directory_handles_deletion_errors(implementation):
    """Test that clear_directory handles deletion errors gracefully."""
    impl_name, module = implementation

    with tempfile.TemporaryDirectory() as parent_dir:
        # Create a directory inside the temporary directory
        test_dir = os.path.join(parent_dir, "test_dir")
        os.makedirs(test_dir)

        # Mock shutil.rmtree and os.rmdir to raise an OSError
        with patch('shutil.rmtree', side_effect=OSError("Mock error")), \
             patch('os.rmdir', side_effect=OSError("Mock error")), \
             patch('builtins.print') as mock_print:

            sig = inspect.signature(module.clear_directory)
            param_names = list(sig.parameters.keys())

            # Skip test if implementation doesn't have a deletion parameter
            if len(param_names) < 2:
                pytest.skip(f"Implementation {impl_name} does not support deletion parameter")

            deletion_param = param_names[1]

            try:
                # Call clear_directory with deletion parameter set to True
                kwargs = {deletion_param: True}
                module.clear_directory(test_dir, **kwargs)

                # Function should print an error message but not crash
                assert mock_print.called, "Function should print an error message"
            except OSError:
                # If OSError was raised, check if it was at least logged
                assert mock_print.called, "Function should print an error before raising"

def test_clear_directory_implementation_behavior(implementation):
    """
    Test that the actual behavior of the implementation matches expected behavior
    by checking the calls to shutil.rmtree and os.remove.
    """
    impl_name, module = implementation

    with tempfile.TemporaryDirectory() as temp_dir:
        # Create some test files and subdirectories
        test_file = os.path.join(temp_dir, "test_file.txt")
        test_subdir = os.path.join(temp_dir, "test_subdir")

        with open(test_file, "w") as f:
            f.write("test content")

        os.makedirs(test_subdir)

        # Mock the relevant functions
        with patch('os.remove', autospec=True) as mock_remove, \
             patch('shutil.rmtree', autospec=True) as mock_rmtree, \
             patch('builtins.print'), \
             patch('os.rmdir', autospec=True) as mock_rmdir:

            # Call clear_directory without deletion
            module.clear_directory(temp_dir)

            # Clear contents - check various implementation approaches
            content_removal_occurred = (
                mock_remove.called or mock_rmtree.called or
                # Count actual rmdir calls excluding potential calls on the dir itself
                sum(1 for call_args in mock_rmdir.call_args_list
                    if call_args[0][0] != temp_dir)
            )
            assert content_removal_occurred, \
                f"Implementation {impl_name} should remove files or directories"

            # Reset mocks
            mock_remove.reset_mock()
            mock_rmtree.reset_mock()
            mock_rmdir.reset_mock()

            # Get the parameter name for deletion
            sig = inspect.signature(module.clear_directory)
            param_names = list(sig.parameters.keys())

            # Skip test if implementation doesn't have a deletion parameter
            if len(param_names) < 2:
                pytest.skip(f"Implementation {impl_name} does not support deletion parameter")

            deletion_param = param_names[1]

            # Call clear_directory with deletion parameter set to True
            kwargs = {deletion_param: True}
            module.clear_directory(temp_dir, **kwargs)

            # Check that directory removal was attempted - either via rmtree or rmdir
            assert mock_rmtree.called or mock_rmdir.called, \
                f"Implementation {impl_name} should attempt to remove the entire directory"

def test_clear_directory_docs_updated(implementation):
    """Test that the docstring for clear_directory has been updated to mention deletion."""
    impl_name, module = implementation

    # Get the docstring
    docstring = module.clear_directory.__doc__ or ""

    # Check if function has a second parameter first
    sig = inspect.signature(module.clear_directory)
    param_names = list(sig.parameters.keys())

    # Skip test if implementation doesn't have a deletion parameter
    if len(param_names) < 2:
        pytest.skip(f"Implementation {impl_name} does not support deletion parameter")

    # Get the parameter name for more accurate testing
    deletion_param = param_names[1]

    # Docstring should mention deletion or related terms
    deletion_terms = ["delet", "remov", "drop"]

    # Check for either the exact parameter name or general deletion terms
    param_mentioned = deletion_param.lower() in docstring.lower()
    terms_mentioned = any(term in docstring.lower() for term in deletion_terms)

    assert param_mentioned or terms_mentioned, \
        f"Implementation {impl_name}'s docstring should mention the deletion capability"

def test_clear_directory_preserves_created_empty_dir(implementation):
    """Test that clear_directory preserves an empty directory it just created."""
    impl_name, module = implementation

    with tempfile.TemporaryDirectory() as parent_dir:
        # Define a non-existent directory path
        nonexistent_dir = os.path.join(parent_dir, "nonexistent_dir")

        # Make sure it doesn't exist
        if os.path.exists(nonexistent_dir):
            shutil.rmtree(nonexistent_dir)

        # Mock print function to avoid output during tests
        with patch('builtins.print'):
            # Call clear_directory on non-existent directory
            module.clear_directory(nonexistent_dir)

            # Directory should be created and empty
            assert os.path.exists(nonexistent_dir)
            assert os.path.isdir(nonexistent_dir)
            assert len(os.listdir(nonexistent_dir)) == 0

def test_clear_directory_handles_readonly_files(implementation):
    """Test clear_directory handles read-only files correctly."""
    impl_name, module = implementation

    with tempfile.TemporaryDirectory() as temp_dir:
        # Create a read-only file
        readonly_file = os.path.join(temp_dir, "readonly.txt")
        with open(readonly_file, "w") as f:
            f.write("readonly content")

        # Make the file read-only (0o444 = r--r--r--)
        os.chmod(readonly_file, 0o444)

        try:
            # Mock print function to avoid output during tests
            with patch('builtins.print'):
                # Call clear_directory
                module.clear_directory(temp_dir)

            # Directory should still exist
            assert os.path.exists(temp_dir)

            # Read-only file should be removed
            assert not os.path.exists(readonly_file)
            assert len(os.listdir(temp_dir)) == 0

        finally:
            # Make sure we restore write permissions if test fails
            if os.path.exists(readonly_file):
                os.chmod(readonly_file, 0o644)

pytest
pytest-mock
transformers
peft
python-dotenv
torch

import pytest
import os
import sys
import json
from typing import Dict, List, Optional, Any

# Import from local test_utils.py in the same directory
from test_utils import TestUtils, TestResultsManager

# Load all implementations in the current sandbox
implementations = TestUtils.load_all_implementations()
test_results = TestResultsManager()

@pytest.fixture(scope="session")
def sandbox_dir():
    """Fixture to provide the sandbox directory path."""
    return os.path.dirname(os.path.abspath(__file__))

@pytest.fixture(scope="session")
def sandbox_name():
    """Fixture to provide the sandbox name."""
    return os.path.basename(os.path.dirname(os.path.abspath(__file__)))

@pytest.fixture(scope="session")
def all_implementations():
    """Fixture to provide all implementations as a dictionary."""
    return implementations

@pytest.fixture(params=list(implementations.items()))
def implementation(request):
    """Fixture to provide each implementation to tests one at a time."""
    return request.param

@pytest.fixture(scope="session")
def results_manager():
    """Fixture to provide access to the test results manager."""
    return test_results

# Hook for collecting test results
@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    """Pytest hook to collect test results."""
    # Execute all other hooks to obtain the report object
    outcome = yield
    rep = outcome.get_result()

    # We're only interested in the call outcome
    if rep.when == "call":
        if hasattr(item, "callspec") and "implementation" in item.callspec.params:
            # Get implementation name and module
            impl_name, _ = item.callspec.params["implementation"]

            # Get test name
            test_name = item.nodeid.split("::")[-1]

            # Record result
            if rep.passed:
                test_results.record_result(impl_name, test_name, True)
            elif rep.failed:
                error_msg = str(rep.longrepr) if rep.longrepr else "Test failed"
                test_results.record_result(impl_name, test_name, False, error_msg)
            elif rep.skipped:
                skip_reason = rep.longrepr[2] if rep.longrepr else "Test skipped"
                test_results.record_skip(impl_name, test_name, skip_reason)

# Hook to save results at the end of testing
@pytest.hookimpl(trylast=True)
def pytest_sessionfinish(session, exitstatus):
    """Save test results at the end of the test session."""
    test_results.save_results()

import os
import sys
import glob
import re
import importlib.util
import traceback
import types
from typing import Dict, List, Optional, Any, Tuple

class TestUtils:
    @staticmethod
    def discover_implementation_files(directory: str = None) -> List[str]:
        """Find all implementation files in the current sandbox directory."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        patterns = [
            r'modified_code\d+\.py',
            r'new_code\d+\.py',
            r'original_code\.py',
            r'implementation\d*\.py'
        ]

        pattern = re.compile('|'.join(f'({p})' for p in patterns))
        implementations = []

        for file_path in glob.glob(os.path.join(directory, '*.py')):
            if pattern.search(os.path.basename(file_path)):
                implementations.append(file_path)

        # Sort files numerically
        def sort_key(path):
            filename = os.path.basename(path)
            match = re.search(r'(\d+)', filename)
            return int(match.group(1)) if match else 0

        return sorted(implementations, key=sort_key)

    @staticmethod
    def create_mock_module(file_path: str, module_name: str, error_info: str) -> types.ModuleType:
        """Create a mock module that contains error information but can still be tested."""
        # Create a new module object
        mock_module = types.ModuleType(module_name)

        # Add basic attributes
        mock_module.__file__ = file_path
        mock_module.__name__ = module_name
        mock_module.__display_name__ = module_name
        mock_module.__error__ = error_info

        # Add a dummy function that can be detected by test functions
        def dummy_function(*args, **kwargs):
            return f"Error in module: {error_info}"

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

    @staticmethod
    def load_module(file_path: str, module_name: Optional[str] = None) -> Any:
        """
        Safely load a module from a file path with proper error handling.
        If the module has errors, return a mock module that can still be tested.
        """
        if module_name is None:
            module_name = os.path.basename(file_path).replace('.py', '')

        # Create a unique module name to avoid conflicts
        sandbox_id = os.path.basename(os.path.dirname(file_path))
        unique_module_name = f"{sandbox_id}_{module_name}"

        try:
            # First, try to read the file to check for syntax errors
            with open(file_path, 'r') as f:
                source_code = f.read()

            # Check for syntax errors by compiling the code
            try:
                compiled = compile(source_code, file_path, 'exec')
            except SyntaxError as e:
                error_msg = f"Syntax error: {str(e)}"
                print(f"Syntax error in {file_path}: {e}")
                print(f"  Line {e.lineno}, column {e.offset}: {e.text}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module spec
            spec = importlib.util.spec_from_file_location(unique_module_name, file_path)
            if spec is None or spec.loader is None:
                error_msg = f"Could not create spec for {file_path}"
                print(f"Error: {error_msg}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module object
            module = importlib.util.module_from_spec(spec)
            sys.modules[unique_module_name] = module

            # Special handling for execution errors
            try:
                # Execute the module code in a safe way
                spec.loader.exec_module(module)
                # Store the original name for reference
                module.__display_name__ = module_name
                return module
            except Exception as e:
                error_msg = f"Runtime error: {str(e)}"
                traceback_str = traceback.format_exc()
                print(f"Error executing module {file_path}: {e}")
                print(traceback_str)

                # Create a partial module that contains what we loaded before the error
                mock_module = TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

                # Copy any attributes that might have been defined before the error
                for attr_name in dir(module):
                    if not attr_name.startswith('__'):
                        try:
                            setattr(mock_module, attr_name, getattr(module, attr_name))
                        except Exception:
                            pass  # Skip attributes that can't be copied

                return mock_module

        except FileNotFoundError as e:
            error_msg = f"File not found: {str(e)}"
            print(f"Error: {error_msg}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)
        except Exception as e:
            error_msg = f"Unexpected error: {str(e)}"
            print(f"Error loading module {file_path}: {e}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

    @classmethod
    def load_all_implementations(cls, directory: str = None) -> Dict[str, Any]:
        """Load all implementation files in the directory, including those with errors."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        implementations = {}

        implementation_files = cls.discover_implementation_files(directory)
        if not implementation_files:
            print("WARNING: No implementation files found. Check your file naming patterns.")

        for file_path in implementation_files:
            module_name = os.path.basename(file_path).replace('.py', '')
            module = cls.load_module(file_path, module_name)

            # Always add the module, even if it has errors
            implementations[module_name] = module

            if hasattr(module, '__error__'):
                print(f"Loaded with errors: {module_name} - {module.__error__}")
            else:
                print(f"Successfully loaded: {module_name}")

        return implementations

class TestResultsManager:
    def __init__(self):
        self.results = {}
        self.sandbox_name = os.path.basename(os.path.dirname(os.path.abspath(__file__)))

    def record_result(self, impl_name: str, test_name: str, passed: bool,
                     error_msg: Optional[str] = None) -> None:
        """Record a test result for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        if passed:
            self.results[impl_name]["passed"] += 1
        else:
            self.results[impl_name]["failed"] += 1
            if error_msg:
                self.results[impl_name]["errors"].append({
                    "test": test_name,
                    "error": error_msg
                })

    def record_skip(self, impl_name: str, test_name: str,
                   reason: Optional[str] = None) -> None:
        """Record a skipped test for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        self.results[impl_name]["skipped"] += 1
        if reason:
            self.results[impl_name]["errors"].append({
                "test": test_name,
                "error": f"SKIPPED: {reason}"
            })

    def get_winner(self) -> Tuple[Optional[int], Dict]:
        """Determine the winner based on test results."""
        winner = None
        max_passed = -1

        for impl_name, results in self.results.items():
            if impl_name == "original_code":
                continue  # Skip original code when determining winner

            if results["passed"] > max_passed:
                max_passed = results["passed"]
                winner = impl_name
            # Break ties by looking at failure count
            elif results["passed"] == max_passed and winner is not None:
                if results["failed"] < self.results[winner]["failed"]:
                    winner = impl_name

        # Convert winner to numeric index if possible
        winner_index = -1
        if winner and re.match(r'modified_code\d+', winner):
            try:
                winner_index = int(re.search(r'(\d+)', winner).group(1))
            except (AttributeError, ValueError):
                pass

        return winner_index, self.results

    def save_results(self, filename: str = "test_results.json") -> None:
        """Save test results to a JSON file."""
        import json

        winner_index, results = self.get_winner()

        # Check if all tests were skipped
        all_skipped = all(
            stats["skipped"] == stats["passed"] + stats["failed"] + stats["skipped"]
            for impl_name, stats in results.items()
            if impl_name != "original_code"
        )

        output = {
            "winner": winner_index,
            "all_skipped": all_skipped,
            "results": {
                name: {
                    "passed": stats["passed"],
                    "failed": stats["failed"],
                    "skipped": stats["skipped"],
                    "total": stats["passed"] + stats["failed"] + stats["skipped"]
                }
                for name, stats in results.items()
                if not name.startswith("_")  # Skip internal items
            }
        }

        with open(filename, "w") as f:
            json.dump(output, f, indent=2)

        print(f"Test results saved to {filename}")

        return output