probing_experiment.py

print("hi from python")
import os
import pickle
import random
import time
import argparse
import math
import json
from tqdm import tqdm
print("did standard imports")

import torch
print("did torch import")
import numpy as np
from sklearn.metrics import *
print("did sklearn imports")

from load import *
print("did load imports")
from config import *
print("did config imports")
from utils import timestamp, seed_all, default_argument_parser, MODEL_N_LAYERS
print("did utils imports")
from make_feature_datasets import prepare_feature_dataset
print("did make_feature_datasets import")
from activations.activation_probing_dataset import make_index_mask, load_activation_probing_dataset
print("did activations imports")

from experiments.activations import *
from experiments.probes import *
from experiments.metrics import *
from experiments.inner_loops import *
print("did experiments imports")


def save_result(exp_cfg, result, inner_loop_name, feature_name):
    save_path = os.path.join(
        os.getenv('RESULTS_DIR', 'results'),
        exp_cfg.experiment_name,
        exp_cfg.model_name,
        exp_cfg.feature_dataset,
        inner_loop_name
    )
    if not os.path.exists(save_path):
        os.makedirs(save_path, exist_ok=True)
        config_dict = vars(exp_cfg)
        if 'dataset_cfg' in config_dict:
            del config_dict['dataset_cfg']
        with open(os.path.join(save_path, 'config.json'), 'w') as f:
            json.dump(config_dict, f)

    model_name = exp_cfg.model_name.replace('.', ',')
    probe_location = exp_cfg.probe_location.replace('.', ',')
    aggregation = exp_cfg.activation_aggregation
    aggregation = aggregation if aggregation is not None else "none"

    save_name = f'{inner_loop_name}.{feature_name}.{model_name}.{probe_location}.{aggregation}.{layer}.p'
    save_file = os.path.join(save_path, save_name)
    pickle.dump(result, open(save_file, 'wb'))


def get_skip_features(exp_cfg, feature_names):
    feature_collection = exp_cfg.feature_dataset.split('.')[0]
    skip_features = set()
    if feature_collection == 'ewt':
        if exp_cfg.feature_subset == 'upos':
            skip_features = set([
                f for f in feature_names if not f.startswith('upos')
            ])
        elif exp_cfg.feature_subset == 'dep':
            skip_features = set([
                f for f in feature_names if not f.startswith('dep')
            ])
        elif exp_cfg.feature_subset == 'morph':
            skip_features = set([
                f for f in feature_names if f.startswith('upos') or f.startswith('dep')
            ])
        elif exp_cfg.feature_subset == 'not-dep':
            skip_features = set([
                f for f in feature_names if f.startswith('dep')
            ])
    elif feature_collection == 'compound_words':
        if exp_cfg.feature_subset:
            features_to_keep = set(exp_cfg.feature_subset.split(','))
            skip_features = set([
                k for k in feature_datasets.keys() if k not in features_to_keep
            ])
    return skip_features


def run_probe_on_layer(exp_cfg, tokenized_dataset, feature_datasets, layer):
    # TODO: add option to compute activations on the fly
    activation_dataset = load_activation_probing_dataset(exp_cfg, layer)
    index_mask = make_index_mask(exp_cfg, tokenized_dataset, feature_datasets)

    print(f'{timestamp()} finished loading activations for layer {layer}')

    # Run the probing experiments
    skip_features = get_skip_features(exp_cfg, feature_datasets.keys())
    results = {inner_loop: {} for inner_loop in exp_cfg.experiment_type}
    for feature_name, feature_data in tqdm(feature_datasets.items()):
        # prepare the feature specific dataset
        if feature_name in skip_features:
            continue
        if feature_data == None:
            print(f'Warning no feature data for {feature_name}')
            continue

        if exp_cfg.activation_aggregation is None:
            # filter the activation dataset to only include the indices required
            feature_indices, feature_classes = feature_data
            all_required_indices = np.where(index_mask.flatten())[0]
            feature_index_mask = np.isin(all_required_indices, feature_indices)
            feature_activation_dataset = activation_dataset[feature_index_mask, :]
        else:
            feature_ix_mask, feature_classes = feature_data
            all_required_seqs = np.where(feature_ix_mask.sum(axis=1) > 0)[0]
            feature_activation_dataset = activation_dataset[all_required_seqs, :]

        # run each of the inner loops on the feature specific dataset
        for inner_loop_name in exp_cfg.experiment_type:
            inner_loop_fn = INNER_LOOP_FNS.get(inner_loop_name, None)
            if inner_loop_fn is None:
                raise ValueError(
                    f'{inner_loop_name} is not a valid experiment type')

            result = inner_loop_fn(
                exp_cfg,
                feature_activation_dataset.astype(np.float32),
                feature_classes
            )
            print(
                f'{timestamp()} | {feature_name} | {inner_loop_name} | {exp_cfg.model_name} | {layer}')

            if exp_cfg.save_features_together:
                results[inner_loop_name][feature_name] = result
            else:
                save_result(exp_cfg, result, inner_loop_name, feature_name)

    if exp_cfg.save_features_together:
        for inner_loop_name in exp_cfg.experiment_type:
            save_result(
                exp_cfg, results[inner_loop_name], inner_loop_name, 'all')


if __name__ == "__main__":
    # see utils.py for args
    parser = default_argument_parser()
    # experiment params
    parser.add_argument(
        '--normalize_activations', action='store_true',
        help='Normalize activations per neuron to have standard deviation 0.1')
    parser.add_argument(
        '--test_set_frac', default=0.3, type=float,
        help='Fraction of dataset to use as test set')
    parser.add_argument(
        '--save_features_together', action='store_true',
        help='Save features together in a single file')
    parser.add_argument(
        '--feature_subset', default='', type=str,
        help='Subset of features to use (functionality determined by feature dataset prepare_dataset())')
    # probe params
    parser.add_argument(
        '--heuristic_feature_selection_method', default='mean_dif', type=str,
        help='Method feature selection (eg, mean_dif, f_stat, mi')
    parser.add_argument(
        '--osp_heuristic_filter_size', default=50, type=int,
        help='Size of initial heuristic feature selection before osp')
    parser.add_argument(
        '--max_k', default=256, type=int,
        help='Max k to use for any inner loop')
    parser.add_argument(
        '--osp_upto_k', default=8, type=int, help='Highest k to use optimal sparse probing')
    parser.add_argument(
        '--gurobi_timeout', default=60, type=int, help='Max time (seconds) to let Gurobi solve')
    parser.add_argument(
        '--gurobi_verbose', default=False, type=bool, help='Print out full gurobi logs')
    parser.add_argument(
        '--min_layer', default=0, type=int, 
        help='Probe will be run on all layers after min_layer.')
    # inner loop specific params
    parser.add_argument(
        '--iterative_pruning_fixed_k', default=5, type=int,
        help='Value of k to hold fixed while implementing iterative pruning')
    parser.add_argument(
        '--iterative_pruning_n_prune_steps', default=10, type=int,
        help='Number of steps to take in iterative pruning')

    args = vars(parser.parse_args())
    feature_dataset_cfg = parse_dataset_args(args['feature_dataset'])
    exp_cfg = ExperimentConfig(args, feature_dataset_cfg)

    seed_all(exp_cfg.seed)

    feature_dataset_info = prepare_feature_dataset(exp_cfg)
    tokenized_dataset, feature_datasets = feature_dataset_info

    print(f'{timestamp()} finished preparing dataset')

    # if model is 'deduped' then it should map to the same entry in MODEL_N_LAYERS
    if args['model'].endswith('-deduped'):
        n_layers = MODEL_N_LAYERS[args['model'][:-8]]
    else:
        n_layers = MODEL_N_LAYERS[args['model']]
    # for parallelization on cluster
    task_id = int(os.getenv('SLURM_ARRAY_TASK_ID', 1))
    task_count = int(os.getenv('SLURM_ARRAY_TASK_COUNT', 1))
    layers = list(range(task_id-1, n_layers, task_count))

    for layer in layers[args['min_layer']:]:
        run_probe_on_layer(
            exp_cfg,
            tokenized_dataset,
            feature_datasets,
            layer
        )