title: TypeError: (-0.7481077572209469+0j) is not JSON serializable in LIT package
body: <p>I am using LIT api to visualise my attention weights.
  Here is the link:  <a href="https://github.com/PAIR-code/lit" rel="nofollow noreferrer">https://github.com/PAIR-code/lit</a></p>
  <p>If I change to <code>PCA</code> in the web interface of output. It throws me this error:</p>
  <pre><code>TypeError: (-0.7481077572209469+0j) is not JSON serializable.
  </code></pre>
  <p>Here is the code:</p>
  <pre><code>from lit_nlp.api import types as lit_types
  from lit_nlp.examples.datasets import glue
  import tensorflow_datasets as tfds
  # https://github.com/PAIR-code/lit/wiki/api.md#adding-models-and-data
  
  import sys
  from absl import app
  from absl import flags
  from absl import logging
  from lit_nlp import dev_server
  from lit_nlp import server_flags
  from lit_nlp.api import model as lit_model
  from lit_nlp.lib import utils
  from lit_nlp.api import dataset as lit_dataset
  from lit_nlp.api import types as lit_types
  from transformers import BertTokenizer, BertForSequenceClassification
  import pandas as pd
  import torch
  import transformers
  from lit_nlp.components import gradient_maps
  
  
  
  
  df = pd.read_excel(&quot;data.xlsx&quot;,sheet_name='master_data')
  
  def load_tfds(*args, do_sort=True, **kw):
      &quot;&quot;&quot;Load from TFDS, with optional sorting.&quot;&quot;&quot;
      # Materialize to NumPy arrays.
      # This also ensures compatibility with TF1.x non-eager mode, which doesn't
      # support direct iteration over a tf.data.Dataset.
  
      # ds = tfds.load('glue/sst2', split='train', shuffle_files=True,download=True)
      ret = df.values.tolist()
      # print(ret)
      # if do_sort:
      #     # Recover original order, as if you loaded from a TSV file.
      #     ret.sort(key=lambda ex: ex['idx'])
      return ret
  
  
  class SST2Data(lit_dataset.Dataset):
      &quot;&quot;&quot;Stanford Sentiment Treebank, binary version (SST-2).
      See https://www.tensorflow.org/datasets/catalog/glue#gluesst2.
      &quot;&quot;&quot;
  
      LABELS = ['0', '1']
  
      def __init__(self, data):
          self._examples = []
          for ex in load_tfds(df):
              self._examples.append({'sentence': ex[0],'label': self.LABELS[ex[1]], })
  
          # print(self._examples)
  
      def spec(self):
          return {
              'sentence': lit_types.TextSegment(),
              'label': lit_types.CategoryLabel(vocab=self.LABELS)
          }
  
  
  
  FLAGS = flags.FLAGS
  
  FLAGS.set_default(&quot;development_demo&quot;, True)
  
  flags.DEFINE_string(
      &quot;model_path&quot;,
      &quot;https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_tiny.tar.gz&quot;,
      &quot;Path to trained model, in standard transformers format, e.g. as &quot;
      &quot;saved by model.save_pretrained() and tokenizer.save_pretrained()&quot;)
  
  
  def _from_pretrained(cls, *args, **kw):
      &quot;&quot;&quot;Load a transformers model in PyTorch, with fallback to TF2/Keras weights.&quot;&quot;&quot;
      try:
          return cls.from_pretrained(*args, **kw)
      except OSError as e:
          logging.warning(&quot;Caught OSError loading model: %s&quot;, e)
          logging.warning(&quot;Re-trying to convert from TensorFlow checkpoint (from_tf=True)&quot;)
          return cls.from_pretrained(*args, from_tf=True, **kw)
  
  
  class SimpleSEModel(lit_model.Model):
      &quot;&quot;&quot;Simple SE classification model.&quot;&quot;&quot;
  
      @property
      def num_layers(self):
          return self.model.config.num_hidden_layers
      @property
      def max_seq_length(self):
          return self.model.config.max_position_embeddings
  
      LABELS = [&quot;0&quot;, &quot;1&quot;]  # negative, positive
  
      def __init__(self, model_name_or_path):
          self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
  
          # This is a just a regular PyTorch model.
          self.model = BertForSequenceClassification.from_pretrained('bert-base-uncased',num_labels=2,output_hidden_states=True,output_attentions=True)
          self.model.eval()
  
      ##
      # LIT API implementation
      def max_minibatch_size(self):
          # This tells lit_model.Model.predict() how to batch inputs to
          # predict_minibatch().
          # Alternately, you can just override predict() and handle batching yourself.
          return 16
  
      def predict_minibatch(self, inputs):
          # Preprocess to ids and masks, and make the input batch.
          encoded_input = self.tokenizer.batch_encode_plus([ex[&quot;sentence&quot;] for ex in inputs],return_tensors=&quot;pt&quot;,add_special_tokens=True,max_length=512,padding=&quot;longest&quot;,truncation=&quot;longest_first&quot;)
  
          # Check and send to cuda (GPU) if available
          if torch.cuda.is_available():
              self.model.cuda()
              for tensor in encoded_input:
                  encoded_input[tensor] = encoded_input[tensor].cuda()
          # Run a forward pass.
          with torch.no_grad():  # remove this if you need gradients.
              out: transformers.modeling_outputs.SequenceClassifierOutput = self.model(**encoded_input)
              unused_attentions = out.attentions
              # print(unused_attentions)
              # print(type(unused_attentions))
  
          # Post-process outputs.
          batched_outputs = {
              &quot;probas&quot;: torch.nn.functional.softmax(out.logits, dim=-1).tolist(),
              &quot;input_ids&quot;: encoded_input[&quot;input_ids&quot;],
              &quot;ntok&quot;: torch.sum(encoded_input[&quot;attention_mask&quot;], dim=1).tolist(),
              &quot;cls_emb&quot;: out.hidden_states[-1][:, 0].tolist(),  # last layer, first token
          }
  
          for i in range(len(unused_attentions)):
              batched_outputs[f&quot;layer_{i:d}_attention&quot;] = unused_attentions[i].detach().numpy()
  
          # unbatched_outputs = utils.unbatch_preds(batched_outputs)
          # Return as NumPy for further processing.
  
          # for k, v in batched_outputs.items():
          #     print(&quot;batched_output&quot;)
          #     print(v)
          #     print(type(v))
          detached_outputs = {k: v for k, v in batched_outputs.items()}
          print(&quot;detached_outputs&quot;)
          print(detached_outputs)
          # Unbatch outputs so we get one record per input example.
          for output in utils.unbatch_preds(detached_outputs):
              ntok = output.pop(&quot;ntok&quot;)
              output[&quot;tokens&quot;] = self.tokenizer.convert_ids_to_tokens(output.pop(&quot;input_ids&quot;)[1:ntok - 1])
              print('output[&quot;tokens&quot;]')
              print(output)
              yield output
  
      def input_spec(self) -&gt; lit_types.Spec:
          return {
              &quot;sentence&quot;: lit_types.TextSegment(),
              &quot;label&quot;: lit_types.CategoryLabel(vocab=self.LABELS, required=False),
              &quot;tokens&quot;: lit_types.Tokens(parent='text', required=False)
          }
  
      def output_spec(self) -&gt; lit_types.Spec:
          spec = {
              &quot;tokens&quot;: lit_types.Tokens(parent=&quot;text&quot;),
              &quot;probas&quot;: lit_types.MulticlassPreds(parent=&quot;label&quot;, vocab=self.LABELS,null_idx=0),
              &quot;cls_emb&quot;: lit_types.Embeddings()}
          for i in range(self.num_layers):
              spec[f&quot;layer_{i:d}_attention&quot;] = lit_types.AttentionHeads(align_in=&quot;tokens&quot;, align_out=&quot;tokens&quot;)
              # print(&quot;spec&quot;)
              # print(spec)
          return spec
  
  
  def get_wsgi_app():
      &quot;&quot;&quot;Returns a LitApp instance for consumption by gunicorn.&quot;&quot;&quot;
      FLAGS.set_default(&quot;server_type&quot;, &quot;external&quot;)
      FLAGS.set_default(&quot;demo_mode&quot;, True)
      # Parse flags without calling app.run(main), to avoid conflict with
      # gunicorn command line flags.
      unused = flags.FLAGS(sys.argv, known_only=True)
      return main(unused)
  
  
  def main(_):
      # Normally path is a directory; if it's an archive file, download and
      # extract to the transformers cache.
      model_path = FLAGS.model_path
      if model_path.endswith(&quot;.tar.gz&quot;):
          model_path = transformers.file_utils.cached_path(model_path, extract_compressed_file=True)
  
      # Load the model we defined above.
      models = {&quot;sst&quot;: SimpleSEModel(model_path)}
      # print('models')
      # print(models)
      # Load SST-2 validation set from TFDS.
      datasets = {&quot;sst_dev&quot;: SST2Data(df)}
  
      # Start the LIT server. See server_flags.py for server options.
      lit_demo = dev_server.Server(models, datasets, **server_flags.get_flags())
      # print('lit_demo')
      # print(lit_demo)
      return lit_demo.serve()
  
  
  if __name__ == &quot;__main__&quot;:
      app.run(main)
  </code></pre>
  <p>Also I am not seeing <code>Gradient L2 norm</code> in model explanation. I am seeing only <code>LIME</code>?</p>