went back a few commits

app.py

@@ -1,7 +1,5 @@
 import gradio as gr
 
-import torch
-
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import numpy as np
 
@@ -57,7 +55,7 @@ if __name__ == "__main__":
         )
         # Important: don't forget to set `normalize_logits=True` to obtain normalized probabilities (i.e. sum(p) = 1)
         transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, normalize_logits=True)
-        transition_proba = np.exp(transition_scores[0])
+        transition_proba = np.exp(transition_scores)
         # We only have scores for the generated tokens, so pop out the prompt tokens
         input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1]
         generated_ids = outputs.sequences[:, input_length:]
@@ -80,7 +78,7 @@ if __name__ == "__main__":
                     break
             highlighted_out.append((token.replace("▁", " "), this_label))
 
-        return highlighted_out, transition_proba
+        return highlighted_out
 
     demo = gr.Blocks()
     with demo: