Spaces:
Sleeping
Sleeping
| 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() | |