mental_health_best_model

mental_health_best_model

LoRA adapter for google/gemma-4-E4B-it fine-tuned jointly on 41 mental-health and emotion text-classification tasks (D01–D22 + GoEmotions one-vs-rest splits) using TRL SFTTrainer with completion-only loss.

The adapter is small (~140 MB) and is meant to be loaded on top of the base model with peft.

Disclaimer — research use only. The outputs are predicted class IDs from a frozen 1-of-K head built into a chat-style prompt; they are not clinical advice and must not be used to diagnose or triage real users.

Training summary

Held-out evaluation (28 % stratified split per task; chat_legacy prompt; 0-shot)

Mean across 41 tasks:

Compared to the same Gemma-4-E4B-it without any fine-tune (zero-shot, same prompt stack):

36/41 tasks improved on BACC; the few regressions are highly imbalanced multiclass tasks (DailyDialog emotion, MultiWD Spiritual) where binary balancing did not apply.

Quick start

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

base_id = "google/gemma-4-E4B-it"
adapter_id = "xxue752/mental_health_best_model"

tokenizer = AutoTokenizer.from_pretrained(adapter_id)
base = AutoModelForCausalLM.from_pretrained(
    base_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
model = PeftModel.from_pretrained(base, adapter_id).eval()

# Single-choice MCQ prompt (mirrors the training format).
system = (
    "You are a careful classifier. Read the input and reply with EXACTLY one "
    "single integer class id from the options listed in the question. No words, "
    "no punctuation, no explanation."
)
question = (
    "Task: single-choice — D01 stress binary.\n"
    "Choose 0 (no stress) or 1 (stress).\n"
    "Input text:\n"
    "I haven't slept in three days because of work and I just can't keep going.\n"
    "Answer:"
)
msgs = [{"role": "system", "content": system}, {"role": "user", "content": question}]
prompt = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
ids = tokenizer(prompt, return_tensors="pt").to(model.device)
out = model.generate(**ids, max_new_tokens=2, do_sample=False)
print(tokenizer.decode(out[0, ids.input_ids.shape[1]:], skip_special_tokens=True))

For the exact 41-task prompt templates and label maps, use the task_core.py module that produced the training data (mirrors mcq_user_message_parts_explicit and make_instruction).

Files

• adapter_config.json, adapter_model.safetensors — PEFT LoRA weights
• tokenizer.json, tokenizer_config.json, chat_template.jinja — tokenizer (Gemma-4 chat template)
• task_sampling_stats.csv — per-task row counts after split / balance / cap
• README.md — this card

Citation

If you use this adapter, please cite the base Gemma model and TRL:

@software{vonwerra2020trl,
  title  = {TRL: Transformer Reinforcement Learning},
  author = {von Werra, Leandro and others},
  url    = {https://github.com/huggingface/trl},
  year   = {2020}
}