Create mental_health_model.py
Browse files- mental_health_model.py +58 -0
mental_health_model.py
ADDED
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import openai
|
2 |
+
from rank_bm25 import BM25Okapi
|
3 |
+
|
4 |
+
class MentalHealthClassifier:
|
5 |
+
def __init__(self, train_data):
|
6 |
+
# Tokenize the training data for BM25
|
7 |
+
self.tokenized_train = [doc.split() for doc in train_data["text"]]
|
8 |
+
self.bm25 = BM25Okapi(self.tokenized_train)
|
9 |
+
self.train_data = train_data
|
10 |
+
|
11 |
+
def classify_text(self, api_key, input_text, k=20):
|
12 |
+
# Set the OpenAI API key
|
13 |
+
openai.api_key = api_key
|
14 |
+
if not openai.api_key:
|
15 |
+
return "Error: OpenAI API key is not set."
|
16 |
+
|
17 |
+
# Tokenize input text
|
18 |
+
tokenized_text = input_text.split()
|
19 |
+
# Get top-k similar examples using BM25
|
20 |
+
scores = self.bm25.get_scores(tokenized_text)
|
21 |
+
top_k_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:k]
|
22 |
+
|
23 |
+
# Build examples for the prompt
|
24 |
+
examples = "\n".join(
|
25 |
+
f"Example {i+1}:\nText: {self.train_data.iloc[idx]['text']}\nClassification: "
|
26 |
+
f"Stress={self.train_data.iloc[idx]['Ground_Truth_Stress']}, "
|
27 |
+
f"Anxiety={self.train_data.iloc[idx]['Ground_Truth_Anxiety']}, "
|
28 |
+
f"Depression={self.train_data.iloc[idx]['Ground_Truth_Depression']}, "
|
29 |
+
f"Other={self.train_data.iloc[idx]['Ground_Truth_Other_binary']}\n"
|
30 |
+
for i, idx in enumerate(top_k_indices)
|
31 |
+
)
|
32 |
+
|
33 |
+
# Construct OpenAI prompt
|
34 |
+
prompt = f"""
|
35 |
+
You are a mental health specialist. Analyze the provided text and classify it into one or more of the following categories: Stress, Anxiety, Depression, or Other.
|
36 |
+
|
37 |
+
Respond with a single category that best matches the content: Stress, Anxiety, Depression, or Other.
|
38 |
+
|
39 |
+
Here is the text to classify:
|
40 |
+
"{input_text}"
|
41 |
+
|
42 |
+
### Examples:
|
43 |
+
{examples}
|
44 |
+
"""
|
45 |
+
|
46 |
+
try:
|
47 |
+
response = openai.ChatCompletion.create(
|
48 |
+
messages=[
|
49 |
+
{"role": "system", "content": "You are a mental health specialist."},
|
50 |
+
{"role": "user", "content": prompt},
|
51 |
+
],
|
52 |
+
model="gpt-4",
|
53 |
+
temperature=0,
|
54 |
+
)
|
55 |
+
content = response.choices[0].message.content.strip()
|
56 |
+
return content # Return the label directly
|
57 |
+
except Exception as e:
|
58 |
+
return f"Error: {e}"
|