import streamlit as st
import utils

from transformers import pipeline
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification

#####################

model_berto='hackathon-somos-nlp-2023/DiagTrast-Berto'
tokenizer_berto = AutoTokenizer.from_pretrained(model_berto)
classifier_berto = pipeline("text-classification", model=model_berto)

model_xml='hackathon-somos-nlp-2023/DiagTrast-Berto'
tokenizer_xml = AutoTokenizer.from_pretrained(model_xml)
classifier_xml = pipeline("text-classification", model=model_xml)

#####################

st.title('Diagnóstico de Trastornos Mentales')

sintomas = st.text_input(label = 'Introduce síntomas',
                         value = 'Nadie me quiere')

pred_berto = classifier_berto(utils.clean_text(sintomas))
pred_xml = classifier_xml(utils.clean_text(sintomas))

st.markdown('Predicción BERTO: ' + pred_berto[0]['label'])
st.markdown('Predicción xml: ' + pred_xml[0]['label'])