import tensorflow as tf
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')
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,
                  padding='max_length',
                  max_length=30,
                  return_tensors='tf')
  y_pred = model(inp)

  return y_pred.logits


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('\n')
  #

  # print predicted class
  i = np.argmax(y_pred_logits)
  return f"Predicted class: {labels[i]} with proba {y_probs[i]:.4f}\n{labels}{y_probs}"

iface = gr.Interface(fn=predict, inputs="text", outputs="text")
iface.launch()