import streamlit as st
from transformers import pipeline


map = {'LABEL_0':'Depression', 'LABEL_1':'Pain', 'LABEL_2':'Anxiety', 'LABEL_3':'Acne', 'LABEL_4':'Birth Control'}

def main():
    st.title("Text Classification App")
    review = st.text_area('Enter the review')

    if st.button:
        if review:
            pipe = pipeline("text-classification", model='ErnestBeckham/gpt-patientconditionclassification', tokenizer='ErnestBeckham/gpt-patientconditionclassification')
            label = pipe(review)
            st.subheader("Patient's Condition:")
            st.write(map[label[0]['label']])

if __name__ == "__main__":
    main()