added avg logprobs

app.py

@@ -80,7 +80,7 @@ if __name__ == "__main__":
                     break
             highlighted_out.append((token.replace("▁", " "), this_label))
 
-        return highlighted_out
+        return highlighted_out, transition_proba
 
     demo = gr.Blocks()
     with demo:
@@ -120,7 +120,7 @@ if __name__ == "__main__":
 
                 avg_logprobs = gr.Textbox(value = torch.mean(transition_proba))
 
-        button.click(get_tokens_and_labels, inputs=prompt, outputs=[highlighted_text, avg_logprobs])
+        button.click(get_tokens_and_labels, inputs=prompt, outputs=[highlighted_text, transition_proba])
 
 
 if __name__ == "__main__":