Spaces:
Running
Running
| import gradio as gr | |
| from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration | |
| import torch | |
| print("Loading MindBridge v3 model...") | |
| model_name = "prats010/mindbridge-mental-health-model" | |
| tokenizer = BlenderbotTokenizer.from_pretrained(model_name) | |
| model = BlenderbotForConditionalGeneration.from_pretrained( | |
| model_name, | |
| torch_dtype=torch.float32 | |
| ) | |
| model.eval() | |
| print("✅ Model loaded!") | |
| def get_max_tokens(user_input): | |
| word_count = len(user_input.split()) | |
| if word_count < 10: | |
| return 220 | |
| else: | |
| return 180 | |
| def get_coping(user_input): | |
| lower = user_input.lower() | |
| if any(w in lower for w in ["anxious", "anxiety", "panic", "worry", "worried"]): | |
| return " Here are some ways to help: try deep breathing exercises, limit caffeine intake, practice grounding techniques like the 5-4-3-2-1 method, and consider speaking with a therapist." | |
| elif any(w in lower for w in ["sad", "depressed", "depression", "hopeless", "empty", "worthless"]): | |
| return " Some things that can help: maintain a daily routine, get sunlight and light exercise, reach out to someone you trust, and consider professional counselling." | |
| elif any(w in lower for w in ["sleep", "insomnia", "tired", "exhausted"]): | |
| return " To improve sleep: avoid screens 1 hour before bed, keep a consistent sleep schedule, try relaxation techniques like body scanning, and limit caffeine after 2pm." | |
| elif any(w in lower for w in ["stress", "overwhelmed", "pressure", "burnout"]): | |
| return " To manage stress: break tasks into smaller steps, take short breaks every 90 minutes, practice mindfulness, and talk to someone about what you are carrying." | |
| elif any(w in lower for w in ["lonely", "alone", "isolated", "nobody"]): | |
| return " To feel more connected: try joining a community or club, reach out to one person today, volunteer, or consider speaking with a counsellor who can offer consistent support." | |
| return "" | |
| def chat(user_input): | |
| if not user_input or not user_input.strip(): | |
| return "Hi, I'm MindBridge. How are you feeling today?" | |
| # Crisis detection | |
| crisis_words = ["suicide", "suicidal", "kill myself", "end my life", "want to die", "self harm", "cutting myself"] | |
| if any(w in user_input.lower() for w in crisis_words): | |
| return "I'm really concerned about what you've shared. Please reach out immediately to iCall at 9152987821 or Vandrevala Foundation at 1860-2662-345. You are not alone and help is available 24/7. [CRISIS]" | |
| inputs = tokenizer( | |
| user_input, | |
| return_tensors="pt", | |
| truncation=True, | |
| max_length=64 | |
| ) | |
| max_tokens = get_max_tokens(user_input) | |
| with torch.no_grad(): | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=max_tokens, | |
| num_beams=4, | |
| temperature=0.8, | |
| do_sample=True, | |
| top_p=0.9, | |
| repetition_penalty=1.3, | |
| ) | |
| response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| response = response + get_coping(user_input) | |
| return response | |
| demo = gr.Interface( | |
| fn=chat, | |
| inputs=gr.Textbox( | |
| label="Your message", | |
| placeholder="How are you feeling today?" | |
| ), | |
| outputs=gr.Textbox(label="Response"), | |
| title="MindBridge AI", | |
| description="Your mental health companion" | |
| ) | |
| demo.launch() |