NeuroMask-v1-165M 🧠
A 165M Parameter Mental Health–Focused Language Model
NeuroMask-v1-165M is a dense, decoder-only Transformer model designed for supportive mental health–related conversational generation.
The model is optimized for English and trained to produce empathetic, counselor-style responses.
📌 Model Summary
- Architecture: Dense Decoder-Only Transformer
- Parameters: 165 Million
- Framework: MaxText
- Training Hardware: TPU v6e-8
- Precision: bfloat16 (Pretraining + Post-training)
- Max Sequence Length: 2048 tokens
- Context Window: 2048 tokens
- Attention (Training): Flash Attention
- Inference Attention (Example): Dot Product
🏗 Training Details
Pretraining
- English language corpus
- General language modeling + mental health–focused text
- Full training performed in bfloat16 precision
- Trained on TPU v6e-8 using MaxText
Post-Training
- Fine-tuned for structured counselor-style dialogue
- Emphasis on empathetic and supportive tone
- Optimized sampling parameters for stable generation
💬 Expected Prompt Format
The model follows a structured dialogue format: User: Counselor:
Example: User: I have anxiety Counselor: Anxiety can be a debilitating condition. I recommend seeking help from a mental health professional who can work with you to develop a treatment plan, which may include cognitive behavioral therapy, medication, and lifestyle changes.User: In addition to therapy, self-care practices like exercise and mindfulness can help manage anxiety symptoms.!User: We can explore different types of therapy, such as cognitive-behavioral therapy, anxiety-provoking techniques, and/or medication. It's important to prioritize your mental health and well-being.!User: If you have questions about the therapy process or are you feeling overwhelmed with your medication, please consider asking for a referral to a medication management specialist.!
🚀 Inference Usage (MaxText)
Run inference using:
python3 -m MaxText.decode maxtext/configs/neuromask.yml \
run_name=inference \
load_parameters_path=~/neuromask-v1-165M/items \
'prompt="User: I have anxiety\nCounselor:"' \
max_target_length=200 \
attention=dot_product \
decode_sampling_strategy=topk \
decode_sampling_top_k=15 \
decode_sampling_temperature=0.7