class_docstr: Ecco's central class. A wrapper around language models. We use it to run the language models
  and collect important data like input saliency and neuron activations.
  
  A LM object is typically not created directly by users,
  it is returned by `ecco.from_pretrained()`.
  
  Usage:
  
  ```python
  import ecco
  
  lm = ecco.from_pretrained('distilgpt2')
  output = lm.generate("Hello computer")
  ```
docstr: Generate tokens in response to an input prompt.
  Works with Language models like GPT2, not masked language models like BERT.
  Args:
      input_str: Input prompt. # TODO: accept batch of input strings
      generate: Number of tokens to generate.
      max_length: max length of sequence (input + output tokens)
      temperature: Adjust the probability distibution of output candidate tokens.
      top_k: Specify top-k tokens to consider in decoding. Only used when do_sample is True.
      top_p: Specify top-p to consider in decoding. Only used when do_sample is True.
      do_sample: Decoding parameter. If set to False, the model always always
          chooses the highest scoring candidate output
          token. This may lead to repetitive text. If set to True, the model considers
          consults top_k and/or top_p to generate more interesting output.
      attribution: List of attribution methods to be calculated. By default, it does not calculate anything.
      beam_size: Beam size to consider while generating
      generate_kwargs: Other arguments to be passed directly to self.model.generate
code: def generate(
      self,
      input_str: str,
      max_length: Optional[int] = 8,
      temperature: Optional[float] = None,
      top_k: Optional[int] = None,
      top_p: Optional[float] = None,
      do_sample: Optional[bool] = False,
      attribution: Optional[List[str]] = [],
      generate: Optional[int] = None,
      beam_size: int = 1,
      **generate_kwargs: Any,
  ):
      """
      Generate tokens in response to an input prompt.
      Works with Language models like GPT2, not masked language models like BERT.
      Args:
          input_str: Input prompt. # TODO: accept batch of input strings
          generate: Number of tokens to generate.
          max_length: max length of sequence (input + output tokens)
          temperature: Adjust the probability distibution of output candidate tokens.
          top_k: Specify top-k tokens to consider in decoding. Only used when do_sample is True.
          top_p: Specify top-p to consider in decoding. Only used when do_sample is True.
          do_sample: Decoding parameter. If set to False, the model always always
              chooses the highest scoring candidate output
              token. This may lead to repetitive text. If set to True, the model considers
              consults top_k and/or top_p to generate more interesting output.
          attribution: List of attribution methods to be calculated. By default, it does not calculate anything.
          beam_size: Beam size to consider while generating
          generate_kwargs: Other arguments to be passed directly to self.model.generate
      """
  
      assert self.model_type in [
          "enc-dec",
          "causal",
      ], f"generate method not supported for model type '{self.model_type}'"
  
      top_k = top_k if top_k is not None else self.model.config.top_k
      top_p = top_p if top_p is not None else self.model.config.top_p
      temperature = (
          temperature
          if temperature is not None
          else self.model.config.temperature
      )
      do_sample = (
          do_sample
          if do_sample is not None
          else self.model.config.task_specific_params.get(
              "text-generation", {}
          ).get("do_sample", False)
      )
  
      pad_token_id = self.model.config.pad_token_id
      eos_token_id = self.model.config.eos_token_id
  
      # We need this as a batch in order to collect activations.
      input_tokenized_info = self.tokenizer(
          input_str, return_tensors="pt"
      )
      input_tokenized_info = self.to(input_tokenized_info)
      input_ids, attention_mask = (
          input_tokenized_info["input_ids"],
          input_tokenized_info["attention_mask"],
      )
      n_input_tokens = len(input_ids[0])
      cur_len = n_input_tokens
  
      if generate is not None:
          max_length = n_input_tokens + generate
  
      if cur_len >= max_length:
          raise ValueError(
              "max_length set to {} while input token has more tokens ({}). Consider increasing max_length".format(
                  max_length, cur_len
              )
          )
  
      # Get decoder input ids
      if (
          self.model_type == "enc-dec"
      ):  # FIXME: only done because causal LMs like GPT-2 have the _prepare_decoder_input_ids_for_generation method but do not use it
          assert len(input_ids.size()) == 2  # will break otherwise
          if version.parse(transformers.__version__) >= version.parse(
              "4.13"
          ):
              # following the code in https://github.com/huggingface/transformers/blob/d0c1aebea467af499331234e7b285a6bf91ea073/tests/generation/test_utils.py#L2099
              model_kwargs = self.model._prepare_encoder_decoder_kwargs_for_generation(
                  input_ids, {}
              )
              (
                  decoder_input_ids,
                  model_kwargs,
              ) = self.model._prepare_decoder_input_ids_for_generation(
                  batch_size=input_ids.shape[0],
                  model_input_name=self.model.main_input_name,
                  model_kwargs=model_kwargs,
                  decoder_start_token_id=self.model.config.decoder_start_token_id,
                  bos_token_id=self.model.config.bos_token_id,
              )
          else:
              decoder_input_ids = (
                  self.model._prepare_decoder_input_ids_for_generation(
                      input_ids, None, None
                  )
              )
      else:
          decoder_input_ids = None
  
      # Print output
      n_printed_tokens = n_input_tokens
      if self.verbose:
          viz_id = self.display_input_sequence(input_ids[0])
  
      # Get model output
      self._remove_hooks()  # deactivate hooks: we will run them for the last model forward only
      output = self.model.generate(
          input_ids=input_ids,
          attention_mask=attention_mask,
          num_beams=beam_size,
          # FIXME: +1 in max_length to account for first start token in decoder, find a better way to do this
          max_length=(generate or max_length - cur_len) + 1
          if self.model_type == "enc-dec"
          else max_length,
          do_sample=do_sample,
          top_p=top_p,
          top_k=top_k,
          temperature=temperature,
          return_dict_in_generate=True,
          output_scores=True,
          **generate_kwargs,
      )
  
      # Get prediction logits for each chosen prediction id
      prediction_logits, prediction_ids = [], []
      if output.__class__.__name__.endswith("EncoderDecoderOutput"):
          prediction_ids, prediction_scores = (
              output.sequences[0][1:],
              output.scores,
          )
  
      elif output.__class__.__name__.endswith("DecoderOnlyOutput"):
          prediction_ids, prediction_scores = (
              output.sequences[0][n_input_tokens:],
              output.scores,
          )
  
      else:
          raise NotImplementedError(
              f"Unexpected output type: {type(output)}"
          )
  
      assert prediction_ids != []
      assert len(prediction_ids) == len(prediction_scores)
      for pred_id, scores in zip(prediction_ids, prediction_scores):
          prediction_logits.append(scores[0][pred_id])
  
      # Analyze each generated token
      self.attributions = defaultdict(list)  # reset attributions dict
      for pred_index, prediction_id in enumerate(prediction_ids):
          # First get encoder/decoder input embeddings
          encoder_input_embeds, _ = self._get_embeddings(input_ids)
  
          # TODO: This is only okay as long as encoder and decoder share the embeddings
          # Should make separate ones for more flexibility
          if decoder_input_ids is not None:
              decoder_input_embeds, _ = self._get_embeddings(
                  decoder_input_ids
              )
          else:
              decoder_input_embeds = None
  
          if (
              pred_index == len(prediction_ids) - 1
          ):  # -1 because we want to catch the inputs for the last generated token
              # attach hooks and run last forward step
              # TODO: collect activation for more than 1 step
              self._attach_hooks(self.model)
              extra_forward_kwargs = {
                  "attention_mask": attention_mask,
                  "decoder_inputs_embeds": decoder_input_embeds,
              }
              forward_kwargs = {
                  "inputs_embeds": encoder_input_embeds,
                  "use_cache": False,
                  "return_dict": True,
                  **{
                      k: v
                      for k, v in extra_forward_kwargs.items()
                      if k
                      in inspect.signature(
                          self.model.forward
                      ).parameters
                  },
              }
              _ = self.model(**forward_kwargs)
  
          # Get primary attributions for produced token
          self._analyze_token(
              encoder_input_embeds=encoder_input_embeds,
              encoder_attention_mask=attention_mask,
              decoder_input_embeds=decoder_input_embeds,
              attribution_flags=attribution,
              prediction_id=prediction_id,
          )
  
          # Recomputing inputs ids, attention mask and decoder input ids
          if decoder_input_ids is not None:
              assert (
                  len(decoder_input_ids.size()) == 2
              )  # will break otherwise
              decoder_input_ids = torch.cat(
                  [
                      decoder_input_ids,
                      torch.tensor(
                          [[prediction_id]],
                          device=decoder_input_ids.device,
                      ),
                  ],
                  dim=-1,
              )
          else:
              input_ids = torch.cat(
                  [
                      input_ids,
                      torch.tensor(
                          [[prediction_id]], device=input_ids.device
                      ),
                  ],
                  dim=-1,
              )
  
              # Recomputing Attention Mask
              if getattr(
                  self.model, "_prepare_attention_mask_for_generation"
              ):
                  assert (
                      len(input_ids.size()) == 2
                  )  # will break otherwise
                  attention_mask = (
                      self.model._prepare_attention_mask_for_generation(
                          input_ids, pad_token_id, eos_token_id
                      )
                  )
                  attention_mask = self.to(attention_mask)
  
          offset = (
              n_input_tokens if decoder_input_ids is not None else 0
          )
          generated_token_ids = (
              decoder_input_ids
              if decoder_input_ids is not None
              else input_ids
          )
  
          # More than one token can be generated at once (e.g., automatic split/pad tokens)
          while (
              len(generated_token_ids[0]) + offset != n_printed_tokens
          ):
              # Display token
              if self.verbose:
                  self.display_token(
                      viz_id,
                      generated_token_ids[0][n_printed_tokens - offset]
                      .cpu()
                      .numpy(),
                      cur_len,
                  )
  
              n_printed_tokens += 1
  
              # Add a zero vector to the attributions vector, if we did not reach the last predicted token
              if (
                  len(generated_token_ids[0]) + offset
                  != n_printed_tokens
              ):
                  for k in self.attributions:
                      self.attributions[k].insert(
                          -1, np.zeros_like(self.attributions[k][-1])
                      )
  
          cur_len += 1
  
      # Get encoder/decoder hidden states
      embedding_states = None
      for attributes in [
          "hidden_states",
          "encoder_hidden_states",
          "decoder_hidden_states",
      ]:
          out_attr = getattr(output, attributes, None)
          if out_attr is not None:
              tokens_hs_list = []
              for token_out_attr in out_attr:
                  hs_list = []
                  for idx, layer_hs in enumerate(token_out_attr):
                      # in Hugging Face Transformers v4, there's an extra index for batch
                      if (
                          len(layer_hs.shape) == 3
                      ):  # If there's a batch dimension, pick the first oen
                          hs = (
                              layer_hs.cpu().detach()[0].unsqueeze(0)
                          )  # Adding a dimension to concat to later
                      # Earlier versions are only 2 dimensional
                      # But also, in v4, for GPT2, all except the last one would have 3 dims, the last layer
                      # would only have two dims
                      else:
                          hs = layer_hs.cpu().detach().unsqueeze(0)
  
                      hs_list.append(hs)
  
                  # First hidden state is the embedding layer, skip it
                  # FIXME: do this in a cleaner way
                  hs_list = torch.cat(hs_list, dim=0)
                  embedding_states = hs_list[0]
                  hidden_states = hs_list[1:]
                  tokens_hs_list.append(hidden_states)
  
              setattr(output, attributes, tokens_hs_list)
  
      # Pass 'hidden_states' to 'decoder_hidden_states'
      if getattr(output, "hidden_states", None) is not None:
          assert (
              getattr(output, "encoder_hidden_states", None) is None
              and getattr(output, "decoder_hidden_states", None) is None
          ), "Not expected to have encoder_hidden_states/decoder_hidden_states with 'hidden_states'"
          setattr(output, "decoder_hidden_states", output.hidden_states)
  
      encoder_hidden_states = getattr(
          output, "encoder_hidden_states", None
      )
      decoder_hidden_states = getattr(
          output,
          "hidden_states",
          getattr(output, "decoder_hidden_states", None),
      )
  
      # Turn activations from dict to a proper array
      activations_dict = self._all_activations_dict
      for layer_type, activations in activations_dict.items():
          self.activations[layer_type] = activations_dict_to_array(
              activations
          )
  
      if decoder_input_ids is not None:
          assert len(decoder_input_ids.size()) == 2
          all_token_ids = torch.cat(
              [input_ids, decoder_input_ids], dim=-1
          )[0]
      else:
          all_token_ids = input_ids[0]
  
      tokens = self.tokenizer.convert_ids_to_tokens(all_token_ids)
      # tokens = []
      # for i in all_token_ids:
      #     token = self.tokenizer.decode([i])
      #     tokens.append(token)
  
      attributions = self.attributions
      attn = getattr(output, "attentions", None)
  
      return OutputSeq(
          **{
              "tokenizer": self.tokenizer,
              "token_ids": all_token_ids.unsqueeze(
                  0
              ),  # Add a batch dimension
              "n_input_tokens": n_input_tokens,
              "output_text": self.tokenizer.decode(all_token_ids),
              "tokens": [tokens],  # Add a batch dimension
              "encoder_hidden_states": encoder_hidden_states,
              "decoder_hidden_states": decoder_hidden_states,
              "embedding_states": embedding_states,
              "attention": attn,
              "attribution": attributions,
              "activations": self.activations,
              "collect_activations_layer_nums": self.collect_activations_layer_nums,
              "lm_head": self.model.lm_head,
              "model_type": self.model_type,
              "device": self.device,
              "config": self.model_config,
          }
      )
  
imports: from collections import defaultdict
  import inspect
  import json
  import os
  import random
  import torch
  import transformers
  from transformers import BatchEncoding
  import ecco
  import numpy
  from IPython import display
  from torch.nn import functional
  from ecco.attribution import compute_primary_attributions_scores
  from ecco.output import OutputSeq
  from typing import Optional, Any, List, Tuple, Dict, Union
  from operator import attrgetter
  import re
  from ecco.util import is_partial_token, strip_tokenizer_prefix
  from packaging import version
label_desc: Use image node_id 5 for calling the LM obj's underlying member method code with example usage: obj.generate(input_str, max_length, temperature, top_k, top_p, do_sample, attribution, generate, beam_size) and returns: OutputSeq