This model is designed to understand users' emotional states, interpret described experiences through the lens of mental health and clinical standards, and provide supportive guidance using a warm, empathetic tone.
This model can be loaded directly using 'transformers'. For optimal performance and role adherence, please strictly follow the parameters below.
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4# Model from HF
5model_id = "hllzmz/medgemma-mentalist"
6
7# Tokenizer
8tokenizer = AutoTokenizer.from_pretrained(model_id)
9
10# Load Model
11# Use load_in_4bit=True for VRAM efficiency (requires bitsandbytes)
12# Or use torch_dtype=torch.bfloat16 for full precision if you have enough VRAM
13model = AutoModelForCausalLM.from_pretrained(
14 model_id,
15 device_map="auto",
16 dtype=torch.bfloat16,
17 load_in_4bit=True
18)
19
20SYSTEM_PROMPT = """
21 You are MedGemma Mentalist, an advanced AI mental health assistant designed to provide empathetic support, scientifically grounded psychoeducation, and guidance.
22 Your goal is to be a bridge to professional help and a source of reliable mental health information.
23 You should NEVER diagnose or stigmatize the user directly.
24 """
25
26def generate_response(user_input):
27 # Construct the conversation history
28 messages = [
29 {"role": "system", "content": SYSTEM_PROMPT},
30 {"role": "user", "content": user_input}
31 ]
32
33 # Apply Chat Template (Gemma Format)
34 input_ids = tokenizer.apply_chat_template(
35 messages,
36 tokenize=True,
37 add_generation_prompt=True,
38 return_tensors="pt"
39 ).to(model.device)
40
41 # Generate Response
42 outputs = model.generate(
43 input_ids=input_ids,
44 max_new_tokens=1024,
45 temperature=0.5,
46 top_p=0.9,
47 repetition_penalty=1.1,
48 )
49
50 # Decode
51 response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
52 return response
53
54# Test Run
55print(generate_response("I feel anxious and sweaty when I am in crowded places."))
To prevent the model from hallucinating or accidentally roleplaying as the client (User), the following generation settings are highly recommended:
Temperature: 0.3 - 0.5 (Lower values ensure the model remains objective and grounded in mental health knowledge).
Repetition Penalty: 1.1 (Prevents the model from getting stuck in loops).