Spaces:
Sleeping
Sleeping
| """Functions to help with searching codes using regex.""" | |
| import pickle | |
| import re | |
| import numpy as np | |
| import torch | |
| from tqdm import tqdm | |
| def load_dataset_cache(cache_base_path): | |
| """Load cache files required for dataset from `cache_base_path`.""" | |
| tokens_str = np.load(cache_base_path + "tokens_str.npy") | |
| tokens_text = np.load(cache_base_path + "tokens_text.npy") | |
| token_byte_pos = np.load(cache_base_path + "token_byte_pos.npy") | |
| return tokens_str, tokens_text, token_byte_pos | |
| def load_code_search_cache(cache_base_path): | |
| """Load cache files required for code search from `cache_base_path`.""" | |
| metrics = np.load(cache_base_path + "metrics.npy", allow_pickle=True).item() | |
| with open(cache_base_path + "cb_acts.pkl", "rb") as f: | |
| cb_acts = pickle.load(f) | |
| with open(cache_base_path + "act_count_ft_tkns.pkl", "rb") as f: | |
| act_count_ft_tkns = pickle.load(f) | |
| return cb_acts, act_count_ft_tkns, metrics | |
| def search_re(re_pattern, tokens_text, at_odd_even=-1): | |
| """Get list of (example_id, token_pos) where re_pattern matches in tokens_text. | |
| Args: | |
| re_pattern: regex pattern to search for. | |
| tokens_text: list of example texts. | |
| at_odd_even: to limit matches to odd or even positions only. | |
| -1 (default): to not limit matches. | |
| 0: to limit matches to odd positions only. | |
| 1: to limit matches to even positions only. | |
| This is useful for the TokFSM dataset when searching for states | |
| since the first token of states are always at even positions. | |
| """ | |
| # TODO: ensure that parentheses are not escaped | |
| assert at_odd_even in [-1, 0, 1], f"Invalid at_odd_even: {at_odd_even}" | |
| if re_pattern.find("(") == -1: | |
| re_pattern = f"({re_pattern})" | |
| res = [ | |
| (i, finditer.span(1)[0]) | |
| for i, text in enumerate(tokens_text) | |
| for finditer in re.finditer(re_pattern, text) | |
| if finditer.span(1)[0] != finditer.span(1)[1] | |
| ] | |
| if at_odd_even != -1: | |
| res = [r for r in res if r[1] % 2 == at_odd_even] | |
| return res | |
| def byte_id_to_token_pos_id(example_byte_id, token_byte_pos): | |
| """Convert byte position (or character position in a text) to its token position. | |
| Used to convert the searched regex span to its token position. | |
| Args: | |
| example_byte_id: tuple of (example_id, byte_id) where byte_id is a | |
| character's position in the text. | |
| token_byte_pos: numpy array of shape (num_examples, seq_len) where | |
| `token_byte_pos[example_id][token_pos]` is the byte position of | |
| the token at `token_pos` in the example with `example_id`. | |
| Returns: | |
| (example_id, token_pos_id) tuple. | |
| """ | |
| example_id, byte_id = example_byte_id | |
| index = np.searchsorted(token_byte_pos[example_id], byte_id, side="right") | |
| return (example_id, index) | |
| def get_code_precision_and_recall(token_pos_ids, codebook_acts, cb_act_counts=None): | |
| """Search for the codes that activate on the given `token_pos_ids`. | |
| Args: | |
| token_pos_ids: list of (example_id, token_pos_id) tuples. | |
| codebook_acts: numpy array of activations of a codebook on a dataset with | |
| shape (num_examples, seq_len, k_codebook). | |
| cb_act_counts: array of shape (num_codes,) where `cb_act_counts[cb_name][code]` | |
| is the number of times the code `code` is activated in the dataset. | |
| Returns: | |
| codes: numpy array of code ids sorted by their precision on the given `token_pos_ids`. | |
| prec: numpy array where `prec[i]` is the precision of the code | |
| `codes[i]` for the given `token_pos_ids`. | |
| recall: numpy array where `recall[i]` is the recall of the code | |
| `codes[i]` for the given `token_pos_ids`. | |
| code_acts: numpy array where `code_acts[i]` is the number of times | |
| the code `codes[i]` is activated in the dataset. | |
| """ | |
| codes = np.array( | |
| [ | |
| codebook_acts[example_id][token_pos_id] | |
| for example_id, token_pos_id in token_pos_ids | |
| ] | |
| ) | |
| codes, counts = np.unique(codes, return_counts=True) | |
| recall = counts / len(token_pos_ids) | |
| idx = recall > 0.01 | |
| codes, counts, recall = codes[idx], counts[idx], recall[idx] | |
| if cb_act_counts is not None: | |
| code_acts = np.array([cb_act_counts[code] for code in codes]) | |
| prec = counts / code_acts | |
| sort_idx = np.argsort(prec)[::-1] | |
| else: | |
| code_acts = np.zeros_like(codes) | |
| prec = np.zeros_like(codes) | |
| sort_idx = np.argsort(recall)[::-1] | |
| codes, prec, recall = codes[sort_idx], prec[sort_idx], recall[sort_idx] | |
| code_acts = code_acts[sort_idx] | |
| return codes, prec, recall, code_acts | |
| def get_neuron_precision_and_recall( | |
| token_pos_ids, recall, neuron_acts_by_ex, neuron_sorted_acts | |
| ): | |
| """Get the neurons with the highest precision and recall for the given `token_pos_ids`. | |
| Args: | |
| token_pos_ids: list of token (example_id, token_pos_id) tuples from a dataset over which | |
| the neurons with the highest precision and recall are to be found. | |
| recall: recall threshold for the neurons (this determines their activation threshold). | |
| neuron_acts_by_ex: numpy array of activations of all the attention and mlp output neurons | |
| on a dataset with shape (num_examples, seq_len, num_layers, 2, dim_size). | |
| The third dimension is 2 because we consider neurons from both: attention and mlp. | |
| neuron_sorted_acts: numpy array of sorted activations of all the attention and mlp output neurons | |
| on a dataset with shape (num_layers, 2, dim_size, num_examples * seq_len). | |
| This should be obtained using the `neuron_acts_by_ex` array by rearranging the first two | |
| dimensions to the last dimensions and then sorting the last dimension. | |
| Returns: | |
| best_prec: highest precision amongst all the neurons for the given `token_pos_ids`. | |
| best_neuron_acts: number of activations of the best neuron for the given `token_pos_ids` | |
| based on the threshold determined by the `recall` argument. | |
| best_neuron_idx: tuple of (layer, is_mlp, neuron_id) where `layer` is the layer number, | |
| `is_mlp` is 0 if the neuron is from attention and 1 if the neuron is from mlp, | |
| and `neuron_id` is the neuron's index in the layer. | |
| """ | |
| if isinstance(neuron_acts_by_ex, torch.Tensor): | |
| neuron_acts_on_pattern = torch.stack( | |
| [ | |
| neuron_acts_by_ex[example_id, token_pos_id] | |
| for example_id, token_pos_id in token_pos_ids | |
| ], | |
| dim=-1, | |
| ) # (layers, 2, dim_size, matches) | |
| neuron_acts_on_pattern = torch.sort(neuron_acts_on_pattern, dim=-1).values | |
| else: | |
| neuron_acts_on_pattern = np.stack( | |
| [ | |
| neuron_acts_by_ex[example_id, token_pos_id] | |
| for example_id, token_pos_id in token_pos_ids | |
| ], | |
| axis=-1, | |
| ) # (layers, 2, dim_size, matches) | |
| neuron_acts_on_pattern.sort(axis=-1) | |
| neuron_acts_on_pattern = torch.from_numpy(neuron_acts_on_pattern) | |
| act_thresh = neuron_acts_on_pattern[ | |
| :, :, :, -int(recall * neuron_acts_on_pattern.shape[-1]) | |
| ] | |
| assert neuron_sorted_acts.shape[:-1] == act_thresh.shape | |
| prec_den = torch.searchsorted(neuron_sorted_acts, act_thresh.unsqueeze(-1)) | |
| prec_den = prec_den.squeeze(-1) | |
| prec_den = neuron_sorted_acts.shape[-1] - prec_den | |
| prec = int(recall * neuron_acts_on_pattern.shape[-1]) / prec_den | |
| assert ( | |
| prec.shape == neuron_acts_on_pattern.shape[:-1] | |
| ), f"{prec.shape} != {neuron_acts_on_pattern.shape[:-1]}" | |
| best_neuron_idx = np.unravel_index(prec.argmax(), prec.shape) | |
| best_prec = prec[best_neuron_idx] | |
| best_neuron_act_thresh = act_thresh[best_neuron_idx].item() | |
| best_neuron_acts = neuron_acts_by_ex[ | |
| :, :, best_neuron_idx[0], best_neuron_idx[1], best_neuron_idx[2] | |
| ] | |
| best_neuron_acts = best_neuron_acts >= best_neuron_act_thresh | |
| best_neuron_acts = np.stack(np.where(best_neuron_acts), axis=-1) | |
| return best_prec, best_neuron_acts, best_neuron_idx | |
| def convert_to_adv_name(name, cb_at, gcb=""): | |
| """Convert layer0_head0 to layer0_attn_preproj_gcb0.""" | |
| if gcb: | |
| layer, head = name.split("_") | |
| return layer + f"_{cb_at}_gcb" + head[4:] | |
| else: | |
| return layer + "_" + cb_at | |
| def convert_to_base_name(name, gcb=""): | |
| """Convert layer0_attn_preproj_gcb0 to layer0_head0.""" | |
| split_name = name.split("_") | |
| layer, head = split_name[0], split_name[-1][3:] | |
| if "gcb" in name: | |
| return layer + "_head" + head | |
| else: | |
| return layer | |
| def get_layer_head_from_base_name(name): | |
| """Convert layer0_head0 to 0, 0.""" | |
| split_name = name.split("_") | |
| layer = int(split_name[0][5:]) | |
| head = None | |
| if len(split_name) > 1: | |
| head = int(split_name[-1][4:]) | |
| return layer, head | |
| def get_layer_head_from_adv_name(name): | |
| """Convert layer0_attn_preproj_gcb0 to 0, 0.""" | |
| base_name = convert_to_base_name(name) | |
| layer, head = get_layer_head_from_base_name(base_name) | |
| return layer, head | |
| def get_codes_from_pattern( | |
| re_pattern, | |
| tokens_text, | |
| token_byte_pos, | |
| cb_acts, | |
| act_count_ft_tkns, | |
| gcb="", | |
| topk=5, | |
| prec_threshold=0.5, | |
| at_odd_even=-1, | |
| ): | |
| """Fetch codes that activate on a given regex pattern. | |
| Retrieves at most `top_k` codes that activate with precision above `prec_threshold`. | |
| Args: | |
| re_pattern: regex pattern to search for. | |
| tokens_text: list of example texts of a dataset. | |
| token_byte_pos: numpy array of shape (num_examples, seq_len) where | |
| `token_byte_pos[example_id][token_pos]` is the byte position of | |
| the token at `token_pos` in the example with `example_id`. | |
| cb_acts: dict of codebook activations. | |
| act_count_ft_tkns: dict over all codebooks of number of token activations on the dataset | |
| gcb: "_gcb" for grouped codebooks and "" for non-grouped codebooks. | |
| topk: maximum number of codes to return per codebook. | |
| prec_threshold: minimum precision required for a code to be returned. | |
| at_odd_even: to limit matches to odd or even positions only. | |
| -1 (default): to not limit matches. | |
| 0: to limit matches to odd positions only. | |
| 1: to limit matches to even positions only. | |
| This is useful for the TokFSM dataset when searching for states | |
| since the first token of states are always at even positions. | |
| Returns: | |
| codebook_wise_codes: dict of codebook name to list of | |
| (code, prec, recall, code_acts) tuples. | |
| re_token_matches: number of tokens that match the regex pattern. | |
| """ | |
| byte_ids = search_re(re_pattern, tokens_text, at_odd_even=at_odd_even) | |
| token_pos_ids = [ | |
| byte_id_to_token_pos_id(ex_byte_id, token_byte_pos) for ex_byte_id in byte_ids | |
| ] | |
| token_pos_ids = np.unique(token_pos_ids, axis=0) | |
| re_token_matches = len(token_pos_ids) | |
| codebook_wise_codes = {} | |
| for cb_name, cb in tqdm(cb_acts.items()): | |
| base_cb_name = convert_to_base_name(cb_name, gcb=gcb) | |
| codes, prec, recall, code_acts = get_code_precision_and_recall( | |
| token_pos_ids, | |
| cb, | |
| cb_act_counts=act_count_ft_tkns[base_cb_name], | |
| ) | |
| idx = np.arange(min(topk, len(codes))) | |
| idx = idx[prec[:topk] > prec_threshold] | |
| codes, prec, recall = codes[idx], prec[idx], recall[idx] | |
| code_acts = code_acts[idx] | |
| codes_pr = list(zip(codes, prec, recall, code_acts)) | |
| codebook_wise_codes[base_cb_name] = codes_pr | |
| return codebook_wise_codes, re_token_matches | |
| def get_neurons_from_pattern( | |
| re_pattern, | |
| tokens_text, | |
| token_byte_pos, | |
| neuron_acts_by_ex, | |
| neuron_sorted_acts, | |
| recall_threshold, | |
| at_odd_even=-1, | |
| ): | |
| """Fetch the highest precision neurons that activate on a given regex pattern. | |
| The activation threshold for the neurons is determined by the `recall_threshold`. | |
| Args: | |
| re_pattern: regex pattern to search for. | |
| tokens_text: list of example texts of a dataset. | |
| token_byte_pos: numpy array of shape (num_examples, seq_len) where | |
| `token_byte_pos[example_id][token_pos]` is the byte position of | |
| the token at `token_pos` in the example with `example_id`. | |
| neuron_acts_by_ex: numpy array of activations of all the attention and mlp output neurons | |
| on a dataset with shape (num_examples, seq_len, num_layers, 2, dim_size). | |
| The third dimension is 2 because we consider neurons from both: attention and mlp. | |
| neuron_sorted_acts: numpy array of sorted activations of all the attention and mlp output neurons | |
| on a dataset with shape (num_layers, 2, dim_size, num_examples * seq_len). | |
| This should be obtained using the `neuron_acts_by_ex` array by rearranging the first two | |
| dimensions to the last dimensions and then sorting the last dimension. | |
| recall_threshold: recall threshold for the neurons (this determines their activation threshold). | |
| at_odd_even: to limit matches to odd or even positions only. | |
| -1 (default): to not limit matches. | |
| 0: to limit matches to odd positions only. | |
| 1: to limit matches to even positions only. | |
| This is useful for the TokFSM dataset when searching for states | |
| since the first token of states are always at even positions. | |
| Returns: | |
| best_prec: highest precision amongst all the neurons for the given `token_pos_ids`. | |
| best_neuron_acts: number of activations of the best neuron for the given `token_pos_ids` | |
| based on the threshold determined by the `recall` argument. | |
| best_neuron_idx: tuple of (layer, is_mlp, neuron_id) where `layer` is the layer number, | |
| `is_mlp` is 0 if the neuron is from attention and 1 if the neuron is from mlp, | |
| and `neuron_id` is the neuron's index in the layer. | |
| re_token_matches: number of tokens that match the regex pattern. | |
| """ | |
| byte_ids = search_re(re_pattern, tokens_text, at_odd_even=at_odd_even) | |
| token_pos_ids = [ | |
| byte_id_to_token_pos_id(ex_byte_id, token_byte_pos) for ex_byte_id in byte_ids | |
| ] | |
| token_pos_ids = np.unique(token_pos_ids, axis=0) | |
| re_token_matches = len(token_pos_ids) | |
| best_prec, best_neuron_acts, best_neuron_idx = get_neuron_precision_and_recall( | |
| token_pos_ids, | |
| recall_threshold, | |
| neuron_acts_by_ex, | |
| neuron_sorted_acts, | |
| ) | |
| return best_prec, best_neuron_acts, best_neuron_idx, re_token_matches | |
| def compare_codes_with_neurons( | |
| best_codes_info, | |
| tokens_text, | |
| token_byte_pos, | |
| neuron_acts_by_ex, | |
| neuron_sorted_acts, | |
| at_odd_even=-1, | |
| ): | |
| """Compare codes with the highest precision neurons on the regex pattern of the code. | |
| Args: | |
| best_codes_info: list of CodeInfo objects. | |
| tokens_text: list of example texts of a dataset. | |
| token_byte_pos: numpy array of shape (num_examples, seq_len) where | |
| `token_byte_pos[example_id][token_pos]` is the byte position of | |
| the token at `token_pos` in the example with `example_id`. | |
| neuron_acts_by_ex: numpy array of activations of all the attention and mlp output neurons | |
| on a dataset with shape (num_examples, seq_len, num_layers, 2, dim_size). | |
| The third dimension is 2 because we consider neurons from both: attention and mlp. | |
| neuron_sorted_acts: numpy array of sorted activations of all the attention and mlp output neurons | |
| on a dataset with shape (num_layers, 2, dim_size, num_examples * seq_len). | |
| This should be obtained using the `neuron_acts_by_ex` array by rearranging the first two | |
| dimensions to the last dimensions and then sorting the last dimension. | |
| at_odd_even: to limit matches to odd or even positions only. | |
| -1 (default): to not limit matches. | |
| 0: to limit matches to odd positions only. | |
| 1: to limit matches to even positions only. | |
| This is useful for the TokFSM dataset when searching for states | |
| since the first token of states are always at even positions. | |
| Returns: | |
| codes_better_than_neurons: fraction of codes that have higher precision than the highest | |
| precision neuron on the regex pattern of the code. | |
| code_best_precs: is an array of the precision of each code in `best_codes_info`. | |
| all_best_prec: is an array of the highest precision neurons on the regex pattern. | |
| """ | |
| assert isinstance(neuron_acts_by_ex, np.ndarray) | |
| ( | |
| neuron_best_prec, | |
| all_best_neuron_acts, | |
| all_best_neuron_idxs, | |
| all_re_token_matches, | |
| ) = zip( | |
| *[ | |
| get_neurons_from_pattern( | |
| code_info.regex, | |
| tokens_text, | |
| token_byte_pos, | |
| neuron_acts_by_ex, | |
| neuron_sorted_acts, | |
| code_info.recall, | |
| at_odd_even=at_odd_even, | |
| ) | |
| for code_info in tqdm(best_codes_info) | |
| ], | |
| strict=True, | |
| ) | |
| neuron_best_prec = np.array(neuron_best_prec) | |
| code_best_precs = np.array([code_info.prec for code_info in best_codes_info]) | |
| codes_better_than_neurons = code_best_precs > neuron_best_prec | |
| return codes_better_than_neurons.mean(), code_best_precs, neuron_best_prec | |