Views
No views yet
1from transformers import AutoTokenizer,AutoModelForCausalLM,pipeline
2import torch
3
4model_id = "kingabzpro/Llama-3.1-8B-Instruct-Mental-Health-Classification"
5
6tokenizer = AutoTokenizer.from_pretrained(model_id)
7
8model = AutoModelForCausalLM.from_pretrained(
9 model_id,
10 return_dict=True,
11 low_cpu_mem_usage=True,
12 torch_dtype=torch.float16,
13 device_map="auto",
14 trust_remote_code=True,
15)
16
17text = "I'm trapped in a storm of emotions that I can't control, and it feels like no one understands the chaos inside me"
18prompt = f"""Classify the text into Normal, Depression, Anxiety, Bipolar, and return the answer as the corresponding mental health disorder label.
19text: {text}
20label: """.strip()
21
22pipe = pipeline(
23 "text-generation",
24 model=model,
25 tokenizer=tokenizer,
26 torch_dtype=torch.float16,
27 device_map="auto",
28)
29
30outputs = pipe(prompt, max_new_tokens=2, do_sample=True, temperature=0.1)
31
32print(outputs[0]["generated_text"].split("label: ")[-1].strip())
33
34# Depression1100%|██████████| 300/300 [03:24<00:00, 1.47it/s]
2
3Accuracy: 0.913
4Accuracy for label Normal: 0.972
5Accuracy for label Depression: 0.913
6Accuracy for label Anxiety: 0.667
7Accuracy for label Bipolar: 0.8001 precision recall f1-score support
2
3 Normal 0.92 0.97 0.95 143
4 Depression 0.93 0.91 0.92 115
5 Anxiety 0.75 0.67 0.71 27
6 Bipolar 1.00 0.80 0.89 15
7
8 accuracy 0.91 300
9 macro avg 0.90 0.84 0.87 300
10weighted avg 0.91 0.91 0.91 3001[[139 3 1 0]
2 [ 5 105 5 0]
3 [ 6 3 18 0]
4 [ 1 2 0 12]]