Spaces:
Sleeping
Sleeping
File size: 3,979 Bytes
2d9f5f4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
MODEL_PATH = "bert_final_model_v1"
LABELS = {
0: "Normal",
1: "Distressed",
2: "Suicidal"
}
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
def predict_mental_health(text):
if not text or not text.strip():
return "Please enter some text to analyze.", {}, ""
inputs = tokenizer(
text,
return_tensors="pt",
truncation=True,
max_length=512,
padding=True
)
with torch.no_grad():
outputs = model(**inputs)
probabilities = torch.softmax(outputs.logits, dim=-1)[0]
predicted_class = torch.argmax(probabilities).item()
prediction = LABELS[predicted_class]
confidence = probabilities[predicted_class].item()
prob_dict = {
LABELS[i]: float(probabilities[i].item())
for i in range(len(LABELS))
}
result_text = (
f"**Prediction:** {prediction}\n\n"
f"**Confidence:** {confidence * 100:.1f}%"
)
return result_text, prob_dict, text
custom_css = """
.gradio-container {
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
}
.primary-btn {
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
}
"""
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
gr.Markdown(
"""
# Mental Health Text Analyzer (BERT Final v1)
AI-powered mental health status detection using a fine-tuned BERT model.
This model classifies text into three categories: **Normal**, **Distressed**, or **Suicidal**.
"""
)
with gr.Row():
with gr.Column():
text_input = gr.Textbox(
label="Enter text to analyze",
placeholder="Type or paste text here...",
lines=5
)
submit_btn = gr.Button("Analyze Text", variant="primary", elem_classes="primary-btn")
gr.Markdown(
"""
### Examples
Try these sample texts to see how the model works.
"""
)
gr.Examples(
examples=[
["I had a wonderful day at the park with my family!"],
["I'm feeling really anxious about my upcoming exam."],
["I feel so hopeless, like nothing will ever get better."],
["Just finished a great workout session, feeling energized!"],
["I can't stop these dark thoughts, everything feels pointless."]
],
inputs=text_input
)
with gr.Column():
result_output = gr.Markdown(label="Result")
probabilities_output = gr.Label(label="Detailed Probabilities", num_top_classes=3)
submit_btn.click(
fn=predict_mental_health,
inputs=text_input,
outputs=[result_output, probabilities_output, text_input]
)
text_input.submit(
fn=predict_mental_health,
inputs=text_input,
outputs=[result_output, probabilities_output, text_input]
)
gr.Markdown(
"""
---
### Important Disclaimer
This tool is for research and educational purposes only. It should not be used as a substitute
for professional mental health care. If you or someone you know is experiencing a mental health crisis,
please contact a mental health professional or crisis helpline immediately.
### Crisis Resources
- National Suicide Prevention Lifeline: 1-800-273-8255
- Crisis Text Line: Text HOME to 741741
- International Association for Suicide Prevention: https://www.iasp.info/resources/Crisis_Centres/
"""
)
if __name__ == "__main__":
demo.launch()
|