Update model.py

model.py

@@ -21,17 +21,14 @@ class ArxivModel:
     def get_idx_class(self, tweet_text, thr=-1.0):
         logits = self.get_logits(tweet_text)
 
-        print(logits)
-
         if thr == -1.0:
-            return [np.argmax(logits)]
+            return [(np.argmax(logits), np.max(logits))]
         else:
             sum_probs = 0.0
             idxs = []
             for p in np.argsort(logits)[::-1]:
                 sum_probs += logits[p]
-                idxs.append(p)
-
+                idxs.append((p, logits[p]))
                 if sum_probs > thr:
                     return idxs