Simplify relative path to mock files

[danbraunai-apollo]

In the tests we have dirty stuff like this:

@pytest.mark.slow
def test_mnist_build_graph():
    mock_config = """
    exp_name: test
    mlp_path: OVERWRITE/IN/MOCK
    batch_size: 256
    seed: 0
    truncation_threshold: 1e-6
    rotate_final_node_layer: false
    n_intervals: 0
    dtype: float32
    node_layers:
        - layers.1
        - layers.2
        - output
    """
    load_config_path = "experiments.mnist_rib_build.run_mnist_rib_build.load_config"

    def mock_load_config_mnist(*args, **kwargs):
        # Load the config as normal but set the mlp_path using a relative path
        config = load_config(*args, **kwargs)
        config.mlp_path = (
            Path(__file__).parent.parent
            / "experiments/train_mnist/sample_checkpoints/lr-0.001_bs-64_2023-08-13_16-23-59/model_epoch_3.pt"
        )
        return config

But, as Nix wrote in his parallelization PR, we can just do:

@pytest.mark.slow
def test_distributed_calc_gives_same_edges():
    rib_dir = str(Path(__file__).parent.parent)
Comment
Lol yep this is way simpler than what I've done in the other files with the mocks, I made an issue here.

    def make_config(name: str, temp_dir: str):
        config_str = f"""
        exp_name: {name}
        seed: 0
        tlens_pretrained: null
        tlens_model_path: {rib_dir}/experiments/train_modular_arithmetic/sample_checkpoints/lr-0.001_bs-10000_norm-None_2023-09-27_18-19-33/model_epoch_60000.pt
        node_layers:
            - ln1.0
            - mlp_in.0
            - unembed
            - output
        dataset:
            source: custom
            name: modular_arithmetic
            return_set: train
        batch_size: 1024
        edge_batch_size: 1024
        truncation_threshold: 1e-6
        rotate_final_node_layer: false
        last_pos_module_type: add_resid1
        n_intervals: 0
        dtype: float64
        eval_type: accuracy
        out_dir: {temp_dir}
        """
        config_path = f"{temp_dir}/{name}.yaml"
        with open(config_path, "w") as f:
            f.write(config_str)
        return config_path

[nix-apollo]

Nicest version of this might be a function that takes a config object and a temp_dir, then converts the config to yaml itself.

[danbraunai-apollo]

Just thought some more on it, I think an even nicer version is just to have the "main" function in all the scripts take a Config object as an argument. Then have the entrypoint to the script be fire.Fire(some_other_function), where some_other_function takes in a filepath to a yaml, converts it to a Config, then calls the main function with that config. In all the tests we'd call main and not some_other_function.

Though we might want to reserve the "main" naming for the highest-level function (the one that does preprocessing and calls the long function in the script).

[nix-apollo]

Could also overload main to take either a path to a config or a config object. I don't know how fire.Fire works with overloaded fns, though.

[nix-apollo]

While I'm refactoring this anyways, do we want to support passing in arbitrary kwargs in the command line to modify the config? e.g. python run_lm_rib_build.py modular_arithmatic.yaml --dtype float64. I'm weakly in favor, and think it's pretty easy with fire.

[danbraunai-apollo]

do we want to support passing in arbitrary kwargs in the command line to modify the config? e.g. python run_lm_rib_build.py modular_arithmatic.yaml --dtype float64. I'm weakly in favor, and think it's pretty easy with fire.

I have quite a strong NO vote for this. I've seen multiple instances of this feature go wrong, where configs don't match up because of the kwargs. It also just makes it harder to share configs around for running experiments

[nix-apollo]

I have quite a strong NO vote for this. I've seen multiple instances of this feature go wrong, where configs don't match up because of the kwargs. It also just makes it harder to share configs around for running experiments

Seems reasonable! Happy to defer to you here.