mavinsao's picture
Update app.py
0d00d6d verified
raw
history blame
1.22 kB
# Use a pipeline as a high-level helper
from transformers import pipeline
import streamlit as st
pipe_1 = pipeline("text-classification", model="mavinsao/mi-roberta-base-finetuned-mental-illness")
pipe_2 = pipeline("text-classification", model="mavinsao/roberta-mental-finetuned")
def ensemble_predict(text):
results_1 = pipe_1(text)
results_2 = pipe_2(text)
# Implement your chosen ensemble strategy here.
# Example with simple averaging:
ensemble_scores = {}
for result in results_1 + results_2:
label = result['label']
score = result['score']
ensemble_scores[label] = ensemble_scores.get(label, 0) + score / 2
predicted_label = max(ensemble_scores, key=ensemble_scores.get)
confidence = ensemble_scores[predicted_label]
return predicted_label, confidence
# Streamlit app
st.title('Mental Illness Prediction')
# Input text area for user input
sentence = st.text_area("Enter the long sentence to predict your mental illness state:")
if st.button('Predict'):
# ... (input validation ... )
predicted_label, confidence = ensemble_predict(sentence)
st.write("Predicted label:", predicted_label)
st.write("Confidence:", confidence)