Project: 019_Do_I_Know_This_Entity_Knowledge_Awareness_and_Hall | Function: get_per_layer_latent_scores (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** SAE latent separation score computation implementing the core algorithm for discovering entity recognition directions across model layers with activation frequency analysis

**KEY TERMS:** {
  "acts_labels_dict": "activation data for known/unknown entities",
  "sae_layers": "l layer indices for SAE analysis",
  "latent_scores": "s^known_{l,j} and s^unknown_{l,j} separation scores",
  "get_features_layers": "computation of f^known_{l,j} frequencies",
  "model_alias": "model identifier for SAE analysis"
}

**CONTEXT CODE:**         queries = load_wikidata_queries(model_alias, free_generation=free_generation, entity_type_and_entity_name_format=entity_type_and_entity_name_format)
        queries = [query for query in queries if query['entity_type'] == entity_type]
        prompts = [q['prompt'] for q in queries]
    elif 'triviaqa' in dataset_name:
        queries = load_triviaqa_queries(model_alias, split=None)
        prompts = [q['prompt'] for q in queries]


    if evaluate_on == 'prompts':
        if 'wikidata' in dataset_name:
            # We first get known and unknown entities
            known_entities, unknown_entities = get_known_unknown_splits(queries, return_type='entity', split=split, consider_refusal_label=False)
            set_entities = known_entities + unknown_entities
            # Get queries of known and unknown entities
            queries = [q for q in queries if q['entity'] in set_entities]
            # Get known (correct) and unknown (error) prompts. That is, prompts with correct or incorrect answers
        known_prompts, unknown_prompts = get_known_unknown_splits(queries, return_type='prompt', split=None, consider_refusal_label=consider_refusal_label)
        print(f"Number of known prompts: {len(known_prompts)}")
        print(f"Number of unknown prompts: {len(unknown_prompts)}")
    elif evaluate_on == 'entities':
        # We split into known and unknown entities
        known_entities, unknown_entities = get_known_unknown_splits(queries, return_type='entity', split=split, consider_refusal_label=consider_refusal_label)
        print(f"Number of known entities: {len(known_entities)}")
        print(f"Number of unknown entities: {len(unknown_entities)}")

    # Get latents for all layers and store them
    if evaluate_on == 'prompts':
        acts_labels_dict = get_acts_labels_dict(model_alias, tokenizer, dataloader, queries, prompts, known_prompts, unknown_prompts, input_data_type=evaluate_on, layers=sae_layers)
        if further_split == True:
            # Update acts_labels_dict[layer]
            # Split errors into errors of known and unknown entities
            for layer in acts_labels_dict.keys():
                labels = acts_labels_dict[layer]['labels']
                prompts = acts_labels_dict[layer]['prompts']
                for i, (prompt, label) in enumerate(zip(prompts, labels)):
                    if label == 1:
                        for entity in known_entities:
                            if entity in prompt:
                                acts_labels_dict[layer]['labels'][i] = 3
                                break

    elif evaluate_on == 'entities':
        acts_labels_dict = get_acts_labels_dict(model_alias, tokenizer, dataloader, queries, prompts, known_entities, unknown_entities, input_data_type=evaluate_on, layers=sae_layers)


    return acts_labels_dict



### Latents scores per layer ###
def get_per_layer_latent_scores(model_alias, tokenizer, n_layers, d_model, sae_layers, save=False, **kwargs):
    """
    Compute latent scores for each layer of the model.

    This function processes the model's activations for a specific entity type,
    computes latent scores using the specified method, and optionally saves the results.

    Args:
        model_alias (str): Alias of the model being analyzed.
        tokenizer: Tokenizer associated with the model.
        n_layers (int): Total number of layers in the model.
        d_model (int): Dimensionality of the model's hidden states.
        sae_layers (list): List of layer indices for which to compute SAE scores.

    """
    raise NotImplementedError("Function implementation has been masked for testing")
    
### Scatter plot latent separation scores ###
def scatter_plot_latent_separation_scores_experiment(model_alias, tokenizer, entity_type, tokens_to_cache, n_layers, testing_layers, d_model, entity_type_and_entity_name_format=False):
    # Parameters
    evaluate_on = 'entities' # keep as 'entities' for computing the scatter plot
    scoring_method = 'absolute_difference' # keep this as 'absolute_difference' for computing the scatter plot
    min_activations = None
    split = 'test'
    save = False # save intermediate results from get_features_layers
    repo_id = model_alias_to_sae_repo_id[model_alias]
   
    # Get features scores for all layers in the training set
    feats_layers = read_layer_features(model_alias, testing_layers, entity_type, scoring_method, tokens_to_cache, evaluate_on)
    train_feats_dict = get_top_k_features(feats_layers, k=None)

    # Get features scores for all layers in the test set
    # TODO: add entity_type to dataset_name
    dataloader = get_dataloader(model_alias, tokens_to_cache, n_layers, d_model, f'wikidata_{entity_type}', batch_size=16)
    test_acts_labels_dict = get_acts_labels_dict_(model_alias, tokenizer, dataloader, testing_layers,
                                                  dataset_name=f'wikidata_{entity_type}', evaluate_on=evaluate_on, split=split,
                                                  free_generation=False, consider_refusal_label=False,
                                                  further_split=False, entity_type_and_entity_name_format=entity_type_and_entity_name_format)
    test_feats_layers = get_features_layers(model_alias, test_acts_labels_dict, testing_layers, SAE_WIDTH, repo_id,
                                            save, scoring_method=scoring_method, min_activations=min_activations,
                                            dataset_name=f'wikidata_{entity_type}')

    test_final_feats_dict = get_top_k_features(test_feats_layers, k=None)

    # we always save scatter plots
    fig = plot_all_features(test_final_feats_dict, train_feats_dict, entity_type, k=10, labels=False)
    os.makedirs('./plots/scatter_plots', exist_ok=True)
    fig.savefig(f'./plots/scatter_plots/{model_alias}_feature_activation_frequencies_smaller_{entity_type}_{str(testing_layers)}.png', transparent=True)
    fig.savefig(f'./plots/scatter_plots/{model_alias}_feature_activation_frequencies_smaller_{entity_type}_{str(testing_layers)}.pdf', transparent=True)
    fig.show()

### Searching for the top general latents ###
def get_general_latents(model_alias, entity_types, testing_layers, tokens_to_cache, evaluate_on, scoring_method, filter_with_pile=False):
    """
    Analyze and identify general latents across different entity types.

    This function processes feature scores for all entity types and computes various statistics to identify
    the most general latents. It calculates mean scores, minimum scores, and mean ranks for each latent
    across all entity types, for both known and unknown entities.

    Args:
        model_alias (str): The alias of the model being analyzed.
        entity_types (list): List of entity types to analyze.
        testing_layers (list): List of layers to analyze.
        tokens_to_cache (str): Type of tokens to cache ('model', 'last_eoi', '?', or 'entity').
        evaluate_on (str): What to evaluate on ('prompts' or 'entities').
        scoring_method (str): Method used for scoring ('absolute_difference' or 't_test').
        filter_with_pile (bool): Whether to filter latents with low frequency in the Pile dataset.

    Returns:
        None. Results are saved as JSON files in the specified directory.

    Side effects:
        - Creates a directory structure based on the input parameters.
        - Saves three JSON files for each known/unknown category:
          1. sorted_scores_min: Minimum scores for each latent across entity types.
          2. sorted_scores_mean: Mean scores for each latent across entity types.
          3. sorted_ranks_mean: Mean ranks for each latent across entity types.

    """
    scores = {}
    ranks = {}
    for known_label in ['known', 'unknown']:

        scores[known_label] = defaultdict(list)
        ranks[known_label] = defaultdict(list)

        for entity_type in entity_types:
            feats_layers = read_layer_features(model_alias, testing_layers, entity_type, scoring_method, tokens_to_cache, evaluate_on)
            train_feats_dict = get_top_k_features(feats_layers, k=None)
            for latent_idx in train_feats_dict[known_label].keys():
                latent = train_feats_dict[known_label][latent_idx]
                full_latent_id = f"L{latent['layer']}F{latent['latent_idx']}"
                scores[known_label][full_latent_id].append(latent['score'])
                ranks[known_label][full_latent_id].append(latent_idx)

        scores_mean = {}
        scores_min = {}
        ranks_mean = {}
        for full_latent_id in scores[known_label].keys():


Complete the function implementation replacing "raise NotImplementedError".