Update app.py

app.py

@@ -65,6 +65,7 @@ if st.button('Submit'):
         encoding = tokenizer(text, return_tensors="pt")
         outputs = model(**encoding)
         predictions = outputs.logits.argmax(-1)
+        number = predictions.cpu().detach().numpy().astype(int)
         probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
    
         fig = plt.figure()
@@ -79,7 +80,7 @@ if st.button('Submit'):
         ax.set_ylabel("Predicted category")
         st.pyplot(fig)
         
-        st.write('Category: {} | Probability: {:.1f}%'.format(labels_plot[predictions.cpu().detach().numpy().astype(int)],(probs_plot[predictions])))
+        st.write('Category: {} | Probability: {:.1f}%'.format(number,(probs_plot[predictions])))
    #     output = genQuestion(option, input)
    #     print(output)
    # st.write(output)