Update app.py

app.py

@@ -1,5 +1,7 @@
 import streamlit as st
+import pandas as pd
 import utils
+import time
 
 from transformers import pipeline
 from transformers import AutoTokenizer
@@ -11,20 +13,64 @@ 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'
+model_xml='hackathon-somos-nlp-2023/DiagTrast-xlm-roberta-base'
 tokenizer_xml = AutoTokenizer.from_pretrained(model_xml)
 classifier_xml = pipeline("text-classification", model=model_xml)
 
 #####################
 
-st.title('Diagnóstico de Trastornos Mentales')
+st.title('🏥 Diagnóstico de Trastornos Mentales')
 
-sintomas = st.text_input(label = 'Introduce síntomas y presiona Enter',
+DemoTab, AboutTab = st.tabs(["Demo", "Acerca de"])
+
+with AboutTab:
+    st.subheader("Motivación")
+    st.markdown(
+        "[Colocar aquí la motivación]."
+    )
+
+    st.subheader("Recursos")
+    st.markdown("""
+        Modelos usados:
+        - [hackathon-somos-nlp-2023/DiagTrast-Berto](https://huggingface.co/hackathon-somos-nlp-2023/DiagTrast-Berto)
+        - [hackathon-somos-nlp-2023/DiagTrast-xlm-roberta-base](https://huggingface.co/hackathon-somos-nlp-2023/DiagTrast-xlm-roberta-base)
+        
+        Dataset:
+        - [hackathon-somos-nlp-2023/DiagTrast](https://huggingface.co/datasets/hackathon-somos-nlp-2023/DiagTrast)
+    """)
+
+    st.subheader("Equipo")
+    st.markdown("""
+        - [Alberto Martín Garrido](https://huggingface.co/Stremie)
+        - [Edgar Mencia]()
+        - [Miguel Ángel Solís Orozco](https://huggingface.co/homosapienssapiens)
+        - [Jose Carlos Vílchez Villegas](https://huggingface.co/JCarlos)
+    """)
+
+with DemoTab:
+    with st.form(key="diagtrast_form"):
+        sintomas = st.text_input(label = 'Introduce síntomas:',
                          value = 'La gente cuando me ve por la calle me mira y halagan mi belleza')
 
-pred_berto = classifier_berto(utils.clean_text(sintomas))
-pred_xml = classifier_xml(utils.clean_text(sintomas))
+        submit_button = st.form_submit_button(label="Clasificar")
+
+    if submit_button and not sintomas:
+        st.warning("⚠️ Debe introducir los síntomas.")
+        
+    elif submit_button:  
+        with st.spinner('Clasificando...'):
+            pred_berto = classifier_berto.predict(utils.clean_text(sintomas))
+            pred_xml = classifier_xml.predict(utils.clean_text(sintomas))
+            
+            df = pd.DataFrame({
+                'Modelo': ["DiagTrast-Berto", "DiagTrast-xlm-roberta-base"],
+                'Diagnóstico': [pred_berto[0]['label'], pred_xml[0]['label']],
+                'Score': [f"{pred_berto[0]['score']:.2%}", f"{pred_xml[0]['score']:.2%}"]
+            })
 
-st.markdown('Predicción BERTO: ' + pred_berto[0]['label'])
-st.markdown('Predicción xml: ' + pred_xml[0]['label'])
+        st.markdown("### Resultados:")
+        st.caption("")
 
+        st.dataframe(df, use_container_width=True)
+        st.caption("")
+        alert = st.success("✅ ¡Hecho!")