class_docstr: An OutputSeq object is the result of running a language model on some input data. It contains not only the output
  sequence of words generated by the model, but also other data collecting during the generation process
  that is useful to analyze the model.
  
  In addition to the data, the object has methods to create plots
  and visualizations of that collected data. These include:
  
  - [layer_predictions()](./#ecco.output.OutputSeq.layer_predictions) <br/>
  Which tokens did the model consider as the best outputs for a specific position in the sequence?
  - [rankings()](./#ecco.output.OutputSeq.rankings) <br/>
  After the model chooses an output token for a specific position, this visual looks back at the ranking
  of this token at each layer of the model when it was generated (layers assign scores to candidate output tokens,
  the higher the "probability" score, the higher the ranking of the token).
  - [rankings_watch()](./#ecco.output.OutputSeq.rankings_watch) <br />
  Shows the rankings of multiple tokens as the model scored them for a single position. For example, if the input is
  "The cat \_\_\_", we use this method to observe how the model ranked the words "is", "are", "was" as candidates
  to fill in the blank.
  - [primary_attributions()](./#ecco.output.OutputSeq.primary_attributions) <br />
  How important was each input token in the selection of calculating the output token?
  
  
  To process neuron activations, OutputSeq has methods to reduce the dimensionality and reveal underlying patterns in
  neuron firings. These are:
  
  - [run_nmf()](./#ecco.output.OutputSeq.run_nmf)
class_docstr_tok: ["An","OutputSeq","object","is","the","result","of","running","a","language","model","on","some","input","data",".","It","contains","not","only","the","output","sequence","of","words","generated","by","the","model",",","but","also","other","data","collecting","during","the","generation","process","that","is","useful","to","analyze","the","model",".","In","addition","to","the","data",",","the","object","has","methods","to","create","plots","and","visualizations","of","that","collected","data",".","These","include",":","-","[","layer_predictions","(",")","]","(",".\/","#","ecco.output.OutputSeq.layer_predictions",")","<","br\/",">","Which","tokens","did","the","model","consider","as","the","best","outputs","for","a","specific","position","in","the","sequence","?","-","[","rankings","(",")","]","(",".\/","#","ecco.output.OutputSeq.rankings",")","<","br\/",">","After","the","model","chooses","an","output","token","for","a","specific","position",",","this","visual","looks","back","at","the","ranking","of","this","token","at","each","layer","of","the","model","when","it","was","generated","(","layers","assign","scores","to","candidate","output","tokens",",","the","higher","the","``","probability","''","score",",","the","higher","the","ranking","of","the","token",")",".","-","[","rankings_watch","(",")","]","(",".\/","#","ecco.output.OutputSeq.rankings_watch",")","<","br","\/",">","Shows","the","rankings","of","multiple","tokens","as","the","model","scored","them","for","a","single","position",".","For","example",",","if","the","input","is","''","The","cat","\\_\\_\\_","''",",","we","use","this","method","to","observe","how","the","model","ranked","the","words","``","is","''",",","``","are","''",",","``","was","''","as","candidates","to","fill","in","the","blank",".","-","[","primary_attributions","(",")","]","(",".\/","#","ecco.output.OutputSeq.primary_attributions",")","<","br","\/",">","How","important","was","each","input","token","in","the","selection","of","calculating","the","output","token","?","To","process","neuron","activations",",","OutputSeq","has","methods","to","reduce","the","dimensionality","and","reveal","underlying","patterns","in","neuron","firings",".","These","are",":","-","[","run_nmf","(",")","]","(",".\/","#","ecco.output.OutputSeq.run_nmf",")"]
docstr: Visualization plotting the topk predicted tokens after each layer (using its hidden state).
  
  Example:
  ![prediction scores](../../img/layer_predictions_ex_london.png)
  
  Args:
      position: The index of the output token to trace
      topk: Number of tokens to show for each layer
      layer: None shows all layers. Can also pass an int with the layer id to show only that layer
docstr_tok: ["Visualization","plotting","the","topk","predicted","tokens","after","each","layer","(","using","its","hidden","state",")",".","Example",":","!","[","prediction","scores","]","(","..","\/","..","\/img\/layer_predictions_ex_london.png",")","Args",":","position",":","The","index","of","the","output","token","to","trace","topk",":","Number","of","tokens","to","show","for","each","layer","layer",":","None","shows","all","layers",".","Can","also","pass","an","int","with","the","layer","id","to","show","only","that","layer"]
code: def layer_predictions(
      self,
      position: int = 1,
      topk: Optional[int] = 10,
      layer: Optional[int] = None,
      **kwargs,
  ):
      """
      Visualization plotting the topk predicted tokens after each layer (using its hidden state).
  
      Example:
      ![prediction scores](../../img/layer_predictions_ex_london.png)
  
      Args:
          position: The index of the output token to trace
          topk: Number of tokens to show for each layer
          layer: None shows all layers. Can also pass an int with the layer id to show only that layer
      """
  
      assert (
          self.model_type != "mlm"
      ), "method not supported for Masked-LMs"
  
      _, dec_hidden_states = self._get_hidden_states()
      assert (
          dec_hidden_states is not None
      ), "decoder hidden states not found"
  
      if position == 0:
          raise ValueError(
              f"'position' is set to 0. There is never a hidden state associated with this position."
              f"Possible values are 1 and above -- the position of the token of interest in the sequence"
          )
  
      if self.model_type in ["enc-dec", "causal"]:
          # The position is relative. By that means, position self.n_input_tokens + 1 is the first generated token
          offset = 1 if self.model_type == "enc-dec" else 0
          new_position = position - offset - self.n_input_tokens
          assert new_position >= 0, (
              f"position={position} not supported, minimum is "
              f"position={self.n_input_tokens + offset} for the first generated token"
          )
          assert new_position < len(dec_hidden_states), (
              f"position={position} not supported, maximum is "
              f"position={len(dec_hidden_states) - 1 + self.n_input_tokens + offset} "
              f"for the last generated token."
          )
          position = new_position
      else:
          raise NotImplemented(
              f"model_type={self.model_type} not supported"
          )
      dec_hidden_states = dec_hidden_states[position][
          :, -1, :
      ]  # only focus on the hidden states for that particular position
  
      if layer is not None:
          # If a layer is specified, choose it only.
          assert dec_hidden_states is not None
          dec_hidden_states = dec_hidden_states[layer].unsqueeze(0)
  
      k = topk
      top_tokens = []
      probs = []
      data = []
  
      # loop through layer levels
      for layer_no, h in enumerate(dec_hidden_states):
          # Use lm_head to project the layer's hidden state to output vocabulary
          logits = self.lm_head(self.to(h))
          softmax = F.softmax(logits, dim=-1)
  
          # softmax dims are (number of words in vocab) - 50257 in GPT2
          sorted_softmax = self.to(torch.argsort(softmax))
  
          # Not currently used. If we're "watching" a specific token, this gets its ranking
          # idx = sorted_softmax.shape[0] - torch.nonzero((sorted_softmax == watch)).flatten()
  
          layer_top_tokens = [
              self.tokenizer.decode(t) for t in sorted_softmax[-k:]
          ][::-1]
          top_tokens.append(layer_top_tokens)
          layer_probs = (
              softmax[sorted_softmax[-k:]].cpu().detach().numpy()[::-1]
          )
          probs.append(layer_probs.tolist())
  
          # Package in output format
          layer_data = []
          for idx, (token, prob) in enumerate(
              zip(layer_top_tokens, layer_probs)
          ):
              layer_num = layer if layer is not None else layer_no
              layer_data.append(
                  {
                      "token": token,
                      "prob": str(prob),
                      "ranking": idx + 1,
                      "layer": layer_num,
                  }
              )
  
          data.append(layer_data)
  
      d.display(
          d.HTML(
              filename=os.path.join(self._path, "html", "setup.html")
          )
      )
      # d.display(d.HTML(filename=os.path.join(self._path, "html", "basic.html")))
  
      js = f"""
           requirejs(['basic', 'ecco'], function(basic, ecco){{
              const viz_id = basic.init()
  
  
              let pred = new ecco.LayerPredictions({{
                  parentDiv: viz_id,
                  data:{json.dumps(data)}
              }})
              pred.init()
           }}, function (err) {{
              console.log(viz_id, err);
          }})"""
      d.display(d.Javascript(js))
  
      if "printJson" in kwargs and kwargs["printJson"]:
          print(data)
          return data
  
imports: import os
  import json
  import ecco
  from IPython import display
  from ecco import util, lm_plots
  import random
  import matplotlib.pyplot
  import numpy
  import torch
  from torch.nn import functional
  from sklearn import decomposition
  from typing import Dict, Optional, List, Tuple, Union
  from ecco.util import strip_tokenizer_prefix, is_partial_token
label_desc: Use image node_id 11 for calling the OutputSeq obj's underlying member method code with example usage: obj.layer_predictions(position, topk, layer) and returns: data