code: def plot_activations(
      tokens,
      activations,
      vmin=0,
      vmax=2,
      height=60,
      width_scale_per_item=1,
      file_prefix="neuron_activation_plot",
  ):
      """Plots a heatmap showing how active each neuron (row) was with each token
      (columns). Neurons with activation less then masking_threashold are masked.
  
      Args:
        tokens: list of the tokens. Note if you're examining activations
        associated with the token as input or as output.
  
  
      """
      n_tokens = activations.shape[-1]
  
      # Activations lower than this threshold will show up as blank
      masking_threshold = 0.01
      masked = np.ma.masked_less(activations, masking_threshold)
      # mask = tensor.shape()
  
      fig = plt.figure(
          figsize=(activations.shape[1], height)
      )  # Width adjusts to number of tokens
  
      sns.set()
      ax = plt.gca()
  
      v = copy.copy(get_cmap("viridis_r"))
      v.set_bad("white")
  
      g = sns.heatmap(
          activations,
          mask=masked.mask,
          cmap=v,
          ax=ax,
          vmin=vmin,
          vmax=vmax,
          cbar=False,
      )
  
      if tokens:
          ax.set_xticklabels(tokens[-n_tokens:], rotation=0)
      ax.tick_params(axis="x", which="major", labelsize=18)
      ax.set_xlabel("\nOutput Token", fontsize=14)
  
      plt.title("FFNN Activations", fontsize=28)
      plt.tick_params(
          axis="x",  # changes apply to the x-axis
          which="both",  # both major and minor ticks are affected
          left=False,  # ticks along the bottom edge are off
          bottom=True,  # ticks along the bottom edge are off
          top=True,  # ticks along thx top edge are off
          labeltop=True,
          labelbottom=True,
      )  # labels along the bottom edge are off
  
      plt.xticks(rotation=-45)
  
      # # Save Figure to file & download
      fig.set_facecolor("w")
      # time_str = int(datetime.datetime.now().strftime("%s")) * 1000
      # file_name = 'distilgpt2_layer_' + str(layer) + '_' + str(time_str) + '.png'
      # plt.savefig(file_name)
      # plt.close(fig)
  
      plt.show()
  
file_path: ecco/src/ecco/lm_plots.py
imports: import seaborn
  import matplotlib.pyplot
  import numpy
  from matplotlib.patches import Rectangle
  from matplotlib.cm import get_cmap
  import copy
  import warnings
  from matplotlib.colors import ListedColormap
  import matplotlib
  from matplotlib.ticker import FuncFormatter
  from typing import Optional, List, Dict
  import sys
label_desc: Use image node_id 1 for calling a global function with example usage: plot_activations(tokens, activations, vmin, vmax, height, width_scale_per_item, file_prefix) without return types