neuromask-v1-165M

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: <message> Counselor: <response>

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