Update app.py

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
 from transformers import pipeline
-import streamlit as st
 import json
 import torch
 from torch.nn import functional as F
@@ -32,18 +31,20 @@ abstract = st.text_input("Abstract", value="Random variable")
 
 def get_logits(title, abstract):
     text = title + "###" + abstract
-    logits = model(tokenizer(text, return_tensors="pt")['input_ids'])['logits']
+    inputs = tokenizer(text, return_tensors="pt")
+    logits = model(**inputs)['logits']
     return logits
 
 def get_ans(logits):
-    ind = torch.argsort(logits, dim=1, descending=True)
-    logits = F.softmax(logits)
+    logits = F.softmax(logits, dim=1)
+    ind = torch.argsort(logits, dim=1, descending=True).flatten()
     cum_sum = 0
     i = 0
-    while cum_sum < 0.95:
-        cum_sum += logits[0][ind[i]]
-        st.write(f"label {ind2label[ind[i]]} with probability {logits[0][ind[i]] * 100}%")
-        i +=1
+    while cum_sum < 0.95 and i < len(ind):
+        idx = ind[i].item()
+        cum_sum += logits[0][idx].item()
+        st.write(f"label: {ind2label.get(str(idx))} with probability: {logits[0][idx].item() * 100:.2f}%")
+        i += 1
 
 if title or abstract:
     logits = get_logits(title, abstract)