Update app.py

app.py

@@ -64,12 +64,10 @@ def AMP(file):
         for index, batch in enumerate(test_dataloader):
             batchs = {k: v for k, v in batch.items()}
             predict = model(batchs)
-            # tsne_plot(output_feature,batchs)
-            out_probability.extend(np.max(np.array(predict.cpu()),axis=1).tolist())
-            # out_probability.extend(np.array(predict.cpu())[:, -1].tolist())
-            test_argmax = np.argmax(predict.cpu(), axis=1).tolist()
+            out_probability = np.max(np.array(predict.cpu()),axis=1)
+            test_argmax = np.argmax(predict.cpu(), axis=1)
     id2str = {0:"non-AMP", 1:"AMP"}
-    return id2str[test_argmax[0]], out_probability
+    return id2str[test_argmax], out_probability
 
 iface = gr.Interface(fn=AMP,
                      inputs="text",