code: def plot_logit_lens(
      tokens,
      softmax_scores,
      predicted_tokens,
      vmin=0,
      vmax=1,
      token_found_mask=None,
      show_input_tokens: bool = False,
      n_input_tokens: int = 0,
  ):
      # print(tokens, softmax_scores, predicted_tokens, token_found_mask)
      start_token = 0
      if not show_input_tokens:
          start_token = n_input_tokens - 1
          if n_input_tokens == 0:
              warnings.warn(
                  "Setting show_input_tokens to True requires supplying n_input_tokens to exlucde inputs from the plot. Defaulting to 0 and showing the input."
              )
  
      fig = plt.figure(figsize=(20, 8))
      fig.set_facecolor("w")
      # Activations lower than this threshold will show up as blank
      # masking_threshold = 0.01
      # masked = np.ma.masked_less(softmax_scores, masking_threshold)
      ax = plt.gca()
  
      v = copy.copy(get_cmap("viridis_r"))
      v.set_bad("white")
  
      g = sns.heatmap(
          softmax_scores[:, start_token:],
          mask=token_found_mask[:, start_token:],
          cmap=v,
          fmt="",
          ax=ax,
          annot=predicted_tokens[:, start_token:],
          vmin=vmin,
          vmax=vmax,
          linewidths=0.5,
          linecolor="#f0f0f0",
          annot_kws={"size": 22},
          cbar_kws={"label": "Probability of token (softmax score)"},
      )
  
      ax.tick_params(axis="x", which="major", labelsize=18)
  
      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=False,  # ticks along thx top edge are off
          labeltop=False,
          labelbottom=True,
      )  # labels along the bottom edge are off
  
      # Output token labels at the bottom
      ax.set_xticklabels(tokens[start_token + 1 :], rotation=-90)
  
      # Layer names label at the left
      ylabels = [
          "Layer {}".format(n) for n in range(softmax_scores.shape[0])
      ]
      ax.set_yticklabels(ylabels, fontsize=18, rotation=0)
  
      # Input token labels at the top
      ax2 = ax.twiny()
      ax2.set_xlim([0, ax.get_xlim()[1]])
      ax2.set_xticks(ax.get_xticks())
      ax2.set_xticklabels(
          tokens[start_token:-1], fontsize=18, rotation=-90
      )
  
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 4 for calling a global function with example usage: plot_logit_lens(tokens, softmax_scores, predicted_tokens, vmin, vmax, token_found_mask, show_input_tokens, n_input_tokens) without return types