|
|
|
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) |
|
|
|
ensemble_scores = {} |
|
for results in [results_1, results_2]: |
|
for result in results: |
|
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 |
|
|
|
|
|
st.title('Mental Illness Prediction') |
|
|
|
|
|
sentence = st.text_area("Enter the long sentence to predict your mental illness state:") |
|
|
|
if st.button('Predict'): |
|
|
|
predicted_label, confidence = ensemble_predict(sentence) |
|
st.write("Predicted label:", predicted_label) |
|
st.write("Confidence:", confidence) |
|
|