from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
import torch
import gradio as gr


#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
labels = ['Not equivalent', "equivalent"]
model_name = "abdulmatinomotoso/paraphrase_detector"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def get_emotion(sentence1, sentence2):
  input_tensor = tokenizer.encode(sentence1, sentence2, return_tensors="pt")
  logits = model(input_tensor).logits

  softmax = torch.nn.Softmax(dim=1)
  probs = softmax(logits)[0]
  probs = probs.cpu().detach().numpy()
  max_index = np.argmax(probs)
  result = labels[max_index]
  return result
  
demo = gr.Interface(get_emotion, inputs=['text', 'text'],
                    outputs="text",
                    title = "PARAPHRASES_DETECTOR -- Detecting if a pair of sentences are equivalent or not")
                    
if __name__ == "__main__":
    demo.launch(debug=True)