Upload mental_health_text_classification.py

mental_health_text_classification.py

@@ -0,0 +1,256 @@
+# -*- coding: utf-8 -*-
+"""Mental Health Text Classification.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/11fDg8hX2svH1yGzRUU8Ji4ooR_4EoVXL
+
+# **1. Install packages**
+"""
+
+!pip install -q --upgrade \
+    transformers==4.51.0 \
+    datasets==3.1.0 \
+    peft==0.13.2 \
+    accelerate==1.0.1 \
+    evaluate \
+    scikit-learn \
+    matplotlib seaborn wordcloud
+
+"""# **2. Imports**"""
+
+import os
+import warnings
+import torch
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, f1_score, classification_report, confusion_matrix
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    Trainer,
+    TrainingArguments,
+    DataCollatorWithPadding,
+)
+from peft import LoraConfig, get_peft_model, TaskType
+from datasets import Dataset
+from huggingface_hub import login
+from google.colab import userdata, files
+
+warnings.filterwarnings("ignore")
+plt.rcParams['figure.figsize'] = (10, 6)
+sns.set(style="whitegrid")
+
+"""# **3. Hugging Face login**"""
+
+try:
+    login(token=userdata.get('HF_TOKEN'))
+    print("Successfully logged in using Colab secret")
+except Exception as e:
+    print(f"Secret login failed: {e}")
+    try:
+        login()
+        print("Interactive login successful")
+    except:
+        print("Login skipped – may hit rate limits")
+
+"""# **4. Download dataset from Kaggle**"""
+
+if not os.path.exists('/root/.kaggle/kaggle.json'):
+    print("Please upload kaggle.json")
+    uploaded = files.upload()
+    if 'kaggle.json' in uploaded:
+        !mkdir -p ~/.kaggle
+        !cp kaggle.json ~/.kaggle/
+        !chmod 600 ~/.kaggle/kaggle.json
+
+!kaggle datasets download -d priyangshumukherjee/mental-health-text-classification-dataset --unzip -p ./data -q
+print("Downloaded files:", os.listdir('./data'))
+
+"""# **5. Load data & prepare labels**"""
+
+TRAIN_PATH = './data/mental_heath_unbanlanced.csv'
+TEST_PATH  = './data/mental_health_combined_test.csv'
+
+df_train = pd.read_csv(TRAIN_PATH)
+df_test  = pd.read_csv(TEST_PATH)
+
+label2id = {'Normal': 0, 'Depression': 1, 'Anxiety': 2, 'Suicidal': 3}
+id2label = {v: k for k, v in label2id.items()}
+
+df_train['label'] = df_train['status'].map(label2id).astype(int)
+df_test['label']  = df_test['status'].map(label2id).astype(int)
+
+df_train = df_train.rename(columns={'label': 'labels'})
+df_test  = df_test.rename(columns={'label': 'labels'})
+
+# Stratified split
+train_df, val_df = train_test_split(
+    df_train,
+    test_size=0.12,
+    stratify=df_train['labels'],
+    random_state=42
+)
+
+train_ds = Dataset.from_pandas(train_df[['text', 'labels']].reset_index(drop=True))
+val_ds   = Dataset.from_pandas(val_df[['text', 'labels']].reset_index(drop=True))
+test_ds  = Dataset.from_pandas(df_test[['text', 'labels']].reset_index(drop=True))
+
+"""# **6. Tokenization**"""
+
+MODEL_NAME = "microsoft/deberta-v3-base"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+def tokenize_function(examples):
+    return tokenizer(
+        examples["text"],
+        truncation=True,
+        max_length=224,
+        padding=False
+    )
+
+tokenized_train = train_ds.map(tokenize_function, batched=True, remove_columns=["text"])
+tokenized_val   = val_ds.map(tokenize_function, batched=True, remove_columns=["text"])
+tokenized_test  = test_ds.map(tokenize_function, batched=True, remove_columns=["text"])
+
+"""# **7. Load model**"""
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+model = AutoModelForSequenceClassification.from_pretrained(
+    MODEL_NAME,
+    num_labels=4,
+    id2label=id2label,
+    label2id=label2id,
+    ignore_mismatched_sizes=True,
+    torch_dtype=torch.float16,
+).to(device)
+
+"""# **8. Apply LoRA**"""
+
+lora_config = LoraConfig(
+    r=16,
+    lora_alpha=32,
+    target_modules=["query_proj", "value_proj"],
+    lora_dropout=0.05,
+    bias="none",
+    task_type=TaskType.SEQ_CLS,
+    modules_to_save=["classifier"]
+)
+
+model = get_peft_model(model, lora_config)
+model.print_trainable_parameters()
+
+for name, param in model.named_parameters():
+    if param.requires_grad:
+        param.data = param.data.float()
+
+print("Trainable parameters")
+
+"""# **9. Metrics**"""
+
+def compute_metrics(eval_pred):
+    logits, labels = eval_pred
+    preds = np.argmax(logits, axis=1)
+    acc = accuracy_score(labels, preds)
+    f1  = f1_score(labels, preds, average="weighted")
+    return {"accuracy": acc, "f1": f1}
+
+"""# **10. TrainingArguments**"""
+
+training_args = TrainingArguments(
+    output_dir                  = "./mental_health_deberta_lora",
+    num_train_epochs            = 4,
+    per_device_train_batch_size = 8,
+    per_device_eval_batch_size  = 16,
+    gradient_accumulation_steps = 2,
+    learning_rate               = 1.5e-4,
+    weight_decay                = 0.01,
+    warmup_ratio                = 0.1,
+    fp16                        = True,
+    eval_strategy               = "epoch",
+    save_strategy               = "epoch",
+    logging_steps               = 100,
+    load_best_model_at_end      = True,
+    metric_for_best_model       = "f1",
+    greater_is_better           = True,
+    report_to                   = "none",
+    optim                       = "adamw_torch",
+    max_grad_norm               = 0.5,
+    lr_scheduler_type           = "cosine",
+    dataloader_num_workers      = 2,
+    remove_unused_columns       = False,
+)
+
+"""# **11. Trainer**"""
+
+trainer = Trainer(
+    model           = model,
+    args            = training_args,
+    train_dataset   = tokenized_train,
+    eval_dataset    = tokenized_val,
+    tokenizer       = tokenizer,
+    data_collator   = DataCollatorWithPadding(tokenizer=tokenizer),
+    compute_metrics = compute_metrics,
+)
+
+"""# **12. training**"""
+
+print("Starting training...")
+trainer.train()
+
+"""# **13. Evaluate & plot**"""
+
+test_results = trainer.evaluate(tokenized_test)
+print("\nTest results:", test_results)
+
+predictions = trainer.predict(tokenized_test)
+preds = np.argmax(predictions.predictions, axis=1)
+true_labels = predictions.label_ids
+
+print("\nClassification Report:\n")
+print(classification_report(true_labels, preds, target_names=list(id2label.values())))
+
+# Confusion Matrix
+cm = confusion_matrix(true_labels, preds)
+sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
+            xticklabels=list(id2label.values()),
+            yticklabels=list(id2label.values()))
+plt.xlabel("Predicted")
+plt.ylabel("True")
+plt.title("Confusion Matrix – Balanced Test Set")
+plt.show()
+
+"""# **14. Save LoRA adapter**"""
+
+print("Merging LoRA weights into base model...")
+merged_model = model.merge_and_unload()
+
+# Optional: Save merged model locally first (for backup)
+merged_model.save_pretrained("./merged_mental_health_deberta")
+tokenizer.save_pretrained("./merged_mental_health_deberta")
+print("Merged model saved locally.")
+
+"""# **15. Push merged model + tokenizer to Hugging Face Hub**"""
+
+repo_id = "OmidSakaki/mental-health-deberta"
+
+print(f"Pushing merged model to: https://huggingface.co/{repo_id}")
+
+merged_model.push_to_hub(
+    repo_id=repo_id,
+    commit_message="Full merged model after LoRA fine-tuning (4-class mental health classification)",
+    safe_serialization=True,
+    private=False
+)
+
+tokenizer.push_to_hub(
+    repo_id=repo_id,
+    commit_message="Tokenizer for merged mental health model"
+)
+
+print("Upload completed! Model is now live at:", f"https://huggingface.co/{repo_id}")