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)
option = st.selectbox(
     'Discourse Type',
     ('Position', 'Concluding Statement', 'Claim', 'Counterclaim' , 'Evidence', 'Lead', 'Position', 'Rebuttal'))
text = st.text_area('Input Here!')

if text:
  SEP = tokenizer1.sep_token
  text_edited = option + SEP + text
  inputs1 = tokenizer1(text_edited , 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.text(out)