Modelpublic

xxue752/mental_health_best_model

sourceHugging Facegemmaupdated 5mo agoView on Hugging Face

0likes5downloads

Model Card

mentalhealthbest_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

Setting	Value
Base model	`google/gemma-4-E4B-it`
Method	LoRA (r=16, α=32, dropout=0.05), bf16
LoRA targets	`q/k/v/o_proj`, `gate/up/down_proj` on the language tower (258 modules)
Loss	TRL SFT, `completion_only_loss=True` (only the answer-token is supervised)
Tasks	41 (`task_core.TASKS[1..41]`), prompt stack `chat_legacy`
Train split	72 % stratified per task; 28 % held out for eval
Per-task cap	1000 train rows
Balance	binary tasks with task_id ≥ 7 AND >1000 train rows downsampled to 1:1 (≈500/500)
Per-task weighting	Menta-style for tasks 1–6 (1.0/1.0/1.2/1.2/1.5/1.5)
Optimizer	AdamW, lr 2e-4, cosine schedule, warmup 0.03, maxgradnorm 1.0
Batch	per-device 1 × grad accum 16 (effective 16)
Epochs	1
Max seq len	8192

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

Mean across 41 tasks:

Metric	Value
ACC	0.826
BACC	0.760
F1	0.763
Macro-AUPRC	0.684
trivial-ACC (always-majority baseline)	0.612

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

Metric	Zero-shot	LoRA finetune	Δ
ACC	0.701	0.826	+0.124
BACC	0.666	0.760	+0.093
F1	0.636	0.763	+0.127
Macro-AUPRC	0.592	0.684	+0.092

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

python

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:

bibtex

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