import streamlit as st
import re
import fasttext

model = fasttext.load_model("fasttext_model.bin")

def preprocess_input(text):
    text = re.sub(r'[^\w\s\']|\n', ' ', text)
    text = re.sub(' +', ' ', text)
    return text.strip().lower()

def classify_transcript(transcript):
    preprocessed_transcript = preprocess_input(transcript)

    prediction = model.predict(preprocessed_transcript)

    predicted_label = prediction[0][0].replace('__label__', '')

    return predicted_label

def main():
    st.title("FASTTEXT MENTAL HEALTH CLASSIFIER")
    st.write("Type 'exit' in the input box below to end the conversation.")

    user_input = st.text_area("Please enter the transcript of the patient:", "")

    if st.button("Classify"):
        if user_input.lower() == 'exit':
            st.stop()
        else:
            predicted_disease = classify_transcript(user_input)
            st.write(f"Based on the transcript, the predicted disease category is: {predicted_disease}")

if __name__ == "__main__":
    main()