import streamlit as st
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer1 = AutoTokenizer.from_pretrained('kitkeat/distilbert-based-uncased-argumentativewriting')
tokenizer2 = AutoTokenizer.from_pretrained('kitkeat/bert-large-uncased-sparse-90-unstructured-pruneofa-argumentativewriting')
tokenizer3 = AutoTokenizer.from_pretrained('kitkeat/deberta-v3-base-argumentativewriting')
model1 = AutoModelForSequenceClassification.from_pretrained('kitkeat/distilbert-based-uncased-argumentativewriting',num_labels=3)
model2 = AutoModelForSequenceClassification.from_pretrained('kitkeat/bert-large-uncased-sparse-90-unstructured-pruneofa-argumentativewriting',num_labels=3)
model3 = AutoModelForSequenceClassification.from_pretrained('kitkeat/deberta-v3-base-argumentativewriting',num_labels=3)
text = st.text_area('Input Here!')

if text:
  inputs1 = tokenizer1(text, padding=True, truncation=True, return_tensors="pt")
  inputs2 = tokenizer2(text, padding=True, truncation=True, return_tensors="pt")
  inputs3 = tokenizer3(text, padding=True, truncation=True, return_tensors="pt")
  
  outputs1 = model1(**inputs1)
  outputs2 = model2(**inputs2)
  outputs3 = model3(**inputs3)
    
  prediction = outputs1.logits.argmax(dim=-1).item()
  # model.config.id2label
  
  if prediction == 0:
      out = 'Adequate'
  elif prediction == 1:
      out = 'Effective'
  elif prediction == 2:
      out = 'Ineffective'
      
  st.json(out)