"""Routes / endpoints related to reading existing data from Azure blob storage.""" from __future__ import annotations import asyncio from typing import Callable, Sequence, TypeVar from fastapi import FastAPI, HTTPException from neuron_explainer.activation_server.explanation_datasets import get_all_explanation_datasets from neuron_explainer.activation_server.load_neurons import ( NodeIdAndDatasets, load_neuron_from_datasets, resolve_neuron_dataset, ) from neuron_explainer.activation_server.neuron_datasets import ( NeuronDatasetMetadata, get_all_neuron_dataset_metadata, ) from neuron_explainer.activations.activations import ActivationRecord from neuron_explainer.activations.attention_utils import get_attended_to_sequence_lengths from neuron_explainer.activations.derived_scalars import DerivedScalarType from neuron_explainer.activations.derived_scalars.tokens import TokenAndRawScalar, TokenAndScalar from neuron_explainer.explanations.explanations import ( ScoredExplanation, load_neuron_explanations_async, ) from neuron_explainer.pydantic import CamelCaseBaseModel, immutable T = TypeVar("T") @immutable class AttributedScoredExplanation(CamelCaseBaseModel): explanation: str score: float | None """ None means that the explanation's score is undefined, maybe because the neuron is never active. """ dataset_name: str @immutable class ExistingExplanationsRequest(CamelCaseBaseModel): dst: DerivedScalarType layer_index: int activation_index: int # Exactly one of explanation_datasets and neuron_dataset must be specified. If # explanation_datasets is specified, only those datasets will be used. If neuron_dataset is # specified, all of its explanation datasets will be used. explanation_datasets: list[str] neuron_dataset: str | None @immutable class TokenAndRawAttentionScalars(CamelCaseBaseModel): # same as TokenAndAttentionScalars, but without any of the postprocessing # applied token: str scalars: list[float] @immutable class TokenAndAttentionScalars(TokenAndRawAttentionScalars): # attention_scalars can be any quantity that is a scalar per attention head # and per token pair. The most straightforward examples are pre- and post-softmax # attention scores; norm of attention head write vector is another option # Causally masked attention results in attention scalars for each pair of # tokens, which read from an earlier token and write to a later token. The # TokenAndAttentionScalars object for the nth token in a sequence contains # a list of n floats, summarizing in some way the values written to the nth # (current) token by reading from all n-1 earlier tokens, plus the current # (nth) token. Because the things read and written are vectors rather than # scalars, there are presumably many ways to summarize them as scalars; # this object contains just one kind of summary. normalized_scalars: list[float] total_scalar_in: float normalized_total_scalar_in: float max_scalar_in: float normalized_max_scalar_in: float total_scalar_out: float normalized_total_scalar_out: float max_scalar_out: float normalized_max_scalar_out: float @immutable class NeuronRecordResponse(CamelCaseBaseModel): dataset: str max_activation: float top_activations: list[list[TokenAndScalar]] random_sample: list[list[TokenAndScalar]] @immutable class AttentionHeadRecordResponse(CamelCaseBaseModel): dataset: str max_attention_activation: float most_positive_token_sequences: list[list[TokenAndAttentionScalars]] random_token_sequences: list[list[TokenAndAttentionScalars]] def zip_tokens_and_activations( activation_records: list[ActivationRecord], ) -> list[list[TokenAndRawScalar]]: sequences = [] for activation_record in activation_records: sequence = [] for token, activation in zip(activation_record.tokens, activation_record.activations): sequence.append( TokenAndRawScalar( token=token, scalar=activation, ) ) sequences.append(sequence) return sequences def convert_activation_records_to_token_and_activation_lists( all_activation_record_lists: list[list[ActivationRecord]], ) -> list[list[list[TokenAndScalar]]]: zipped_tokens_and_raw_activations = [ zip_tokens_and_activations(activation_record_list) for activation_record_list in all_activation_record_lists ] return normalize_token_scalars(zipped_tokens_and_raw_activations) def zip_tokens_and_attention_activations( activation_records: list[ActivationRecord], ) -> list[list[TokenAndRawAttentionScalars]]: """ This function takes a list of activation records and returns a list of lists of TokenAndRawAttentionScalars, one outer list element per sequence, and one inner list element per sequence token. There are multiple attention activations per sequence token (one per token in the "attended to sequence" for that sequence token). Because the activation record holds the activations in a flattened form, where (sequence token index, attended to sequence token index) are combined into a single index, we need to reconstruct the nested list structure, inferring it from the number of flattened activations and the number of sequence tokens. """ sequences: list[list[TokenAndRawAttentionScalars]] = [] for activation_record in activation_records: sequence = convert_token_sequence_to_token_and_raw_attention_scalars( activation_record.tokens, activation_record.activations ) sequences.append(sequence) return sequences def convert_token_sequence_to_token_and_raw_attention_scalars( token_strings: list[str], activations: list[float] ) -> list[TokenAndRawAttentionScalars]: num_sequence_tokens = len(token_strings) num_activations = len(activations) attended_to_sequence_lengths_indexed_by_sequence_token = get_attended_to_sequence_lengths( num_sequence_tokens, num_activations ) attention_activations_list_indexed_by_sequence_token = [] start = 0 for ( attended_to_sequence_length_for_sequence_token ) in attended_to_sequence_lengths_indexed_by_sequence_token: end = start + attended_to_sequence_length_for_sequence_token attention_activations_list_indexed_by_sequence_token.append(activations[start:end]) start = end assert len(attention_activations_list_indexed_by_sequence_token) == len(token_strings) sequence = [] for token_string, activations in zip( token_strings, attention_activations_list_indexed_by_sequence_token ): sequence.append( TokenAndRawAttentionScalars( token=token_string, scalars=activations, ) ) return sequence def convert_activation_records_to_token_and_attention_activations_lists( all_activation_record_lists: list[list[ActivationRecord]], ) -> list[list[list[TokenAndAttentionScalars]]]: zipped_tokens_and_raw_attention_activations = [ zip_tokens_and_attention_activations(activation_record_list) for activation_record_list in all_activation_record_lists ] return normalize_attention_token_scalars(zipped_tokens_and_raw_attention_activations) def define_read_routes(app: FastAPI) -> None: @app.post( "/existing_explanations", response_model=list[AttributedScoredExplanation], tags=["read"] ) async def existing_explanations( request: ExistingExplanationsRequest, ) -> list[AttributedScoredExplanation]: def convert_scored_explanation( scored_explanation: ScoredExplanation, explanation_dataset: str ) -> AttributedScoredExplanation: return AttributedScoredExplanation( explanation=scored_explanation.explanation, score=scored_explanation.get_preferred_score(), dataset_name=explanation_dataset.split("/")[-1], ) async def load_and_convert_explanations( explanation_dataset: str, ) -> list[AttributedScoredExplanation]: neuron_simulation_results = await load_neuron_explanations_async( explanation_dataset, request.layer_index, request.activation_index ) if neuron_simulation_results is None: return [] else: return [ convert_scored_explanation(scored_explanation, explanation_dataset) for scored_explanation in neuron_simulation_results.scored_explanations if scored_explanation.explanation is not None ] if not ((len(request.explanation_datasets) > 0) ^ (request.neuron_dataset is not None)): raise HTTPException( status_code=400, detail="Exactly one of explanation_datasets and neuron_dataset must be specified.", ) if len(request.explanation_datasets) > 0: explanation_datasets = request.explanation_datasets else: assert request.neuron_dataset is not None # Redundant assert; mypy needs this. neuron_dataset = resolve_neuron_dataset(request.neuron_dataset, request.dst) explanation_datasets = await get_all_explanation_datasets(neuron_dataset) tasks = [load_and_convert_explanations(dataset) for dataset in explanation_datasets] scored_explanation_lists = await asyncio.gather(*tasks) # Flatten the list of lists. return [item for sublist in scored_explanation_lists for item in sublist] @app.post("/neuron_record", response_model=NeuronRecordResponse, tags=["read"]) async def neuron_record(request: NodeIdAndDatasets) -> NeuronRecordResponse: dataset_path, neuron_record = await load_neuron_from_datasets(request) top_activations, random_sample = convert_activation_records_to_token_and_activation_lists( [ neuron_record.most_positive_activation_records, neuron_record.random_sample, ] ) return NeuronRecordResponse( dataset=dataset_path, max_activation=neuron_record.max_activation, top_activations=top_activations, random_sample=random_sample, ) @app.post("/attention_head_record", response_model=AttentionHeadRecordResponse, tags=["read"]) async def attention_head_record(request: NodeIdAndDatasets) -> AttentionHeadRecordResponse: dataset_path, neuron_record = await load_neuron_from_datasets(request) ( most_positive_token_sequences, random_token_sequences, ) = convert_activation_records_to_token_and_attention_activations_lists( [ neuron_record.most_positive_activation_records, neuron_record.random_sample, ] ) return AttentionHeadRecordResponse( dataset=dataset_path, max_attention_activation=neuron_record.max_activation, most_positive_token_sequences=most_positive_token_sequences, random_token_sequences=random_token_sequences, ) @app.post( "/neuron_datasets_metadata", response_model=list[NeuronDatasetMetadata], tags=["read"] ) def neuron_datasets_metadata() -> list[NeuronDatasetMetadata]: return get_all_neuron_dataset_metadata() def flatten(list_of_lists: list[list[T]]) -> list[T]: return [item for sublist in list_of_lists for item in sublist] def normalize_token_scalars( list_of_sequence_lists: list[list[list[TokenAndRawScalar]]], ) -> list[list[list[TokenAndScalar]]]: """The input ist a list of lists of lists of TokenAndRawScalar objects; the outer list is indexed by provenance, e.g. top or random token sequences; the middle list is indexed by token sequence; the inner list is indexed by token within a sequence. flatten() collapses the outermost level of nesting. We first flatten the whole thing, to compute the max activation across all tokens in all sequences of all provenances. We floor at 0, so if all activations are negative max_activation is 0. Then we step through each token in each sequence in each provenance, and divide by that max activation, and floor at 0, to get a normalized activation with ceiling at 1 and floor at 0. This normalized_activation goes in a TokenAndScalar object, which is in a nested list structure parallel to the input structure.""" indexed_by_token: list[TokenAndRawScalar] = flatten(flatten(list_of_sequence_lists)) indexed_by_token = [ TokenAndRawScalar(token=d.token, scalar=max(d.scalar, 0)) for d in indexed_by_token ] max_scalar = max([d.scalar for d in indexed_by_token]) or 0 neuron_scale = create_scale_linear(max_scalar) return [ [ [ TokenAndScalar( token=d.token, scalar=d.scalar, normalized_scalar=neuron_scale(d.scalar), ) for d in sequence ] for sequence in list_of_sequences ] for list_of_sequences in list_of_sequence_lists ] def create_scale_linear(value: float) -> Callable[[float], float]: return lambda x: max((x - 0) / (value - 1e-5), 0) def normalize_attention_token_scalars( list_of_sequence_lists: list[list[list[TokenAndRawAttentionScalars]]], ) -> list[list[list[TokenAndAttentionScalars]]]: """The high level setup is analogous to normalize_token_scalars, but using TokenAndRawAttentionScalars objects, which have a list of floats associated to each token. There are several kinds of normalizations applied: 1. Normalizing all scalars in all provenances, in all sequences, in all tokens, all scalars within a token's scalars, to the same scale so that they are in the range [0, 1]. 2. Summarizing the lists of scalars per token in some way to get a single scalar per token, and then normalizing those scalars to the same scale, as described in 1.and a. one way to summarize is to compute a summary statistic on the "scalars in", which are the list of floats associated to each token. b. the other way is to compute a summary statistic on the "scalars out", which are the ith entry in each list of of floats that contains an ith entry. i. one summary statistic used is the "total", or sum of all scalars in or out ii. another summary statistic used is the "max", or maximum of all scalars in or out These normalized and summarized scalars are then stored in a TokenAndAttentionScalars object, which is returned in a nested list structure parallel to the input structure. """ indexed_by_token_sequence = flatten(list_of_sequence_lists) indexed_by_token = flatten(indexed_by_token_sequence) indexed_by_token = [ TokenAndRawAttentionScalars(token=d.token, scalars=[max(a, 0) for a in d.scalars]) for d in indexed_by_token ] max_scalar = max([a for d in indexed_by_token for a in d.scalars]) or 0 total_scalar_scale = create_scale_linear(max_scalar) def compute_summary_of_scalar_in( attn_token_sequence_list: list[TokenAndRawAttentionScalars], operation: Callable[[Sequence[float | None]], float], ) -> list[float]: return [operation(d.scalars) for d in attn_token_sequence_list] def _get_entry_if_available(scalars: Sequence[float], index: int) -> float | None: # if index is out of bounds, return None; otherwise return the entry at that index return scalars[index] if index < len(scalars) else None def _sum_non_none(scalars: Sequence[float | None]) -> float: # sum all non-None entries in scalars return sum([a for a in scalars if a is not None]) def _max_non_none(scalars: Sequence[float | None]) -> float: # return the max of all non-None entries in scalars return max([a for a in scalars if a is not None]) def compute_summary_of_scalar_out( attn_token_sequence_list: list[TokenAndRawAttentionScalars], operation: Callable[[Sequence[float | None]], float], ) -> list[float]: for i in range(len(attn_token_sequence_list)): # the attended to sequence length at token i is at most i+1 assert len(attn_token_sequence_list[i].scalars) <= i + 1, ( i, len(attn_token_sequence_list[i].scalars), ) scalar_out = [ operation( [ _get_entry_if_available(attn_token_sequence_list[i].scalars, j) for i in range(j, len(attn_token_sequence_list)) ] ) for j in range(len(attn_token_sequence_list)) ] return scalar_out def compute_total_scalar_in( attn_token_sequence_list: list[TokenAndRawAttentionScalars], ) -> list[float]: return compute_summary_of_scalar_in(attn_token_sequence_list, _sum_non_none) def compute_max_scalar_in( attn_token_sequence_list: list[TokenAndRawAttentionScalars], ) -> list[float]: return compute_summary_of_scalar_in(attn_token_sequence_list, _max_non_none) def compute_total_scalar_out( attn_token_sequence_list: list[TokenAndRawAttentionScalars], ) -> list[float]: return compute_summary_of_scalar_out(attn_token_sequence_list, _sum_non_none) def compute_max_scalar_out( attn_token_sequence_list: list[TokenAndRawAttentionScalars], ) -> list[float]: return compute_summary_of_scalar_out(attn_token_sequence_list, _max_non_none) def compute_scalar_summary_and_scale( compute_scalar_summary_function: Callable[[list[TokenAndRawAttentionScalars]], list[float]], list_of_sequence_lists: list[list[list[TokenAndRawAttentionScalars]]], ) -> tuple[list[list[list[float]]], Callable[[float], float]]: scalar_indexed_by_token_sequence_list: list[list[list[float]]] = [ [compute_scalar_summary_function(sequence) for sequence in sequence_list] for sequence_list in list_of_sequence_lists ] scalar_indexed_by_token: list[float] = flatten( flatten(scalar_indexed_by_token_sequence_list) ) scalar_indexed_by_token = [max(a, 0) for a in scalar_indexed_by_token] scalar_scale: Callable[[float], float] = create_scale_linear( max(scalar_indexed_by_token) or 0 ) return scalar_indexed_by_token_sequence_list, scalar_scale ( total_scalar_in_indexed_by_token_sequence_list, total_scalar_in_scale, ) = compute_scalar_summary_and_scale(compute_total_scalar_in, list_of_sequence_lists) ( max_scalar_in_indexed_by_token_sequence_list, max_scalar_in_scale, ) = compute_scalar_summary_and_scale(compute_max_scalar_in, list_of_sequence_lists) ( total_scalar_out_indexed_by_token_sequence_list, total_scalar_out_scale, ) = compute_scalar_summary_and_scale(compute_total_scalar_out, list_of_sequence_lists) ( max_scalar_out_indexed_by_token_sequence_list, max_scalar_out_scale, ) = compute_scalar_summary_and_scale(compute_max_scalar_out, list_of_sequence_lists) return [ [ [ TokenAndAttentionScalars( token=d.token, scalars=d.scalars, normalized_scalars=[total_scalar_scale(a) for a in d.scalars], total_scalar_in=total_scalar_in_indexed_by_token_sequence_list[seq_list_idx][ seq_idx ][token_idx], normalized_total_scalar_in=total_scalar_in_scale( total_scalar_in_indexed_by_token_sequence_list[seq_list_idx][seq_idx][ token_idx ] ), max_scalar_in=max_scalar_in_indexed_by_token_sequence_list[seq_list_idx][ seq_idx ][token_idx], normalized_max_scalar_in=max_scalar_in_scale( max_scalar_in_indexed_by_token_sequence_list[seq_list_idx][seq_idx][ token_idx ] ), total_scalar_out=total_scalar_out_indexed_by_token_sequence_list[seq_list_idx][ seq_idx ][token_idx], normalized_total_scalar_out=total_scalar_out_scale( total_scalar_out_indexed_by_token_sequence_list[seq_list_idx][seq_idx][ token_idx ] ), max_scalar_out=max_scalar_out_indexed_by_token_sequence_list[seq_list_idx][ seq_idx ][token_idx], normalized_max_scalar_out=max_scalar_out_scale( max_scalar_out_indexed_by_token_sequence_list[seq_list_idx][seq_idx][ token_idx ] ), ) for token_idx, d in enumerate(sequence) ] for seq_idx, sequence in enumerate(list_of_sequences) ] for seq_list_idx, list_of_sequences in enumerate(list_of_sequence_lists) ]