import gradio as gr
import tensorflow as tf
import text_hammer as th
from transformers import DistilBertTokenizer, TFDistilBertForSequenceClassification

# Initialize the tokenizer and model (assuming you've already fine-tuned it)
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
model = TFDistilBertForSequenceClassification.from_pretrained("Elegbede-Distilbert_FInetuned_For_Text_Classification")
# Define a function to make predictions
def predict(texts):

  preprocessed_texts = []
  for text in texts:
    text = str(text).lower()
    text = th.cont_exp(text)  # Apply cont_exp function
    text = th.remove_special_chars(text)  # Remove special characters
    text = th.remove_accented_chars(text)  # Remove accented characters
    preprocessed_texts.append(text)

  # Join the list of preprocessed texts back into a single string
  preprocessed_text = ' '.join(preprocessed_texts)

  # Tokenize and preprocess the new text
  new_encodings = tokenizer(preprocessed_text, truncation=True, padding=True, max_length=70, return_tensors='tf')

  # Make predictions
  new_predictions = model(new_encodings)
  new_labels_pred = tf.argmax(new_predictions.logits, axis=1)

  labels_dict = {0: 'Sadness 😭', 1: 'Joy 😂', 2: 'Love 😍', 3: 'Anger 😠', 4: 'Fear 😨', 5: 'Surprise 😲'}

  # Assuming 'new_labels_pred' contains the predicted class index
  predicted_emotion = labels_dict[new_labels_pred[0].numpy()]
  return predicted_emotion

iface = gr.Interface(
  fn=predict,
  inputs="text",
  outputs="text",
  title="Emotion Classifier",
  description="Predict the emotion from text using a fine-tuned DistilBERT model ",
)
# Launch the interfac
iface.launch()