Update app.py

app.py

@@ -3,17 +3,12 @@ from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
 import gradio as gr
 import numpy as np
 from scipy.special import softmax
+from transformers import TFAutoModelForSequenceClassification
 
 tokenizer = AutoTokenizer.from_pretrained('roberta-base')
-
-from transformers import TFAutoModelForSequenceClassification
 access_token = "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"
-
 model = TFAutoModelForSequenceClassification.from_pretrained("ilan541/ssid_classification", use_auth_token=access_token)
 
-
-
-
 def get_probs(text):
   inp = tokenizer(text,
                   truncation=True,
@@ -29,20 +24,24 @@ def predict(your_text):
   
   y_pred_logits = get_probs(your_text)
   labels = ['Bus', 'Portable Device', 'Stationary Device', 'Vehicle']
+  num_labels = len(labels)
+
+  # # print logits
+  # for i in range(num_labels):
+  #   print(f"logit - {labels[i]}: {y_pred_logits[:,i][0]:.4f}")
+
+  # print('\n')
+
+  # # print probas
+  y_probs = softmax(y_pred_logits)[0]
+  # for i in range(num_labels):
+  #   print(f"prob - {labels[i]}: {y_probs[i]:.4f}")
   
-  # print logits
-  for label, logit in zip(labels, y_pred_logits):
-    print(f"logit - {label}: {logit}")
-  
-  # print probas
-  y_probs = softmax(y_pred_logits)
-  for label, prob in zip(labels, y_probs):
-    print(f"prob - {label}: {prob}")
-  
+  # print('\n')
+
   # print predicted class
   i = np.argmax(y_pred_logits)
-  print(f"Predicted class: {labels[i]} with proba {y_probs[i]}")
-  
+  return f"Predicted class: {labels[i]} with proba {y_probs[i]:.4f}"
 
 iface = gr.Interface(fn=predict, inputs="text", outputs="text")
-iface.launch()
+iface.launch()