feat: add module 2 emotion classifier

README.md

@@ -31,6 +31,48 @@ reports/module_1_language_detection/
 
 The UI returns the detected language, confidence, and whether the prediction passed the confidence threshold.
 
+## Module 2: Emotion Classification
+
+The emotion classifier uses a fine-tuned transformer:
+
+- Base model: `distilbert-base-uncased`
+- Dataset: `dair-ai/emotion`
+- Labels: sadness, joy, love, anger, fear, surprise
+- Training target: run on Colab T4 using `notebooks/module_2_emotion_training.ipynb`
+
+DistilBERT is used because it keeps most of BERT's language understanding while being smaller and faster, which makes it a better fit for a student project that needs GPU training but practical local inference.
+
+After training in Colab, the notebook saves:
+
+```text
+src/models/saved_emotion_model/
+reports/module_2_emotion_classification/
+```
+
+The local inference class returns the predicted emotion, confidence, and a simple word-occlusion explanation showing which words most affected the predicted emotion.
+
+### Module Integration Plan
+
+The final chatbot will analyze each user message in this order:
+
+```text
+User message -> Language Detection -> Emotion Classification -> Intent Classification -> RAG/direct response
+```
+
+Module 1 decides the language for routing and response language. Module 2 adds emotional context so later response generation can be gentler for sadness/fear/anger and more direct for neutral informational requests. Crisis handling should still be implemented as a separate safety route later, not inferred from emotion alone.
+
+Run emotion inference after exporting the trained model:
+
+```bash
+.\.venv\Scripts\python.exe src\models\emotion_classifier.py "I feel anxious and overwhelmed" --explain
+```
+
+Run the emotion UI:
+
+```bash
+.\.venv\Scripts\python.exe src\models\emotion_detector_ui.py
+```
+
 ### Install Dependencies
 
 ```bash

notebooks/module_2_emotion_training.ipynb

@@ -0,0 +1,285 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "id": "371f3178",
+      "metadata": {},
+      "source": [
+        "# Module 2: Emotion Classification\n",
+        "\n",
+        "This notebook fine-tunes a transformer emotion classifier for the mental health chatbot. It is designed for Google Colab with a T4 GPU, while the project repo keeps the reusable inference and explanation code locally."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "57bcd49f",
+      "metadata": {},
+      "source": [
+        "## Why DistilBERT?\n",
+        "\n",
+        "We use `distilbert-base-uncased` because it is transformer-based, lighter than BERT, fast enough for Colab T4 training, and practical for local inference. For this project, that balance is better than a larger model that may be harder to deploy or explain during assessment.\n",
+        "\n",
+        "Dataset: `dair-ai/emotion` with six labels: sadness, joy, love, anger, fear, and surprise."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "d22909ef",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!nvidia-smi"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "19fba046",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "%pip -q install -U transformers datasets accelerate evaluate scikit-learn pandas"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "4c57f7d2",
+      "metadata": {},
+      "source": [
+        "## Colab Repo Setup\n",
+        "\n",
+        "Run this notebook from the cloned project folder so the trained model and reports are saved back into the same structure used locally."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "bfb00d86",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# If you opened the notebook outside the repo in Colab, uncomment these lines once:\n",
+        "# !git clone https://github.com/MarwanZaineldeen/Mental-Health-Chatbot.git\n",
+        "# %cd Mental-Health-Chatbot"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "64279088",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import sys\n",
+        "from inspect import signature\n",
+        "from pathlib import Path\n",
+        "\n",
+        "import numpy as np\n",
+        "import pandas as pd\n",
+        "from datasets import load_dataset\n",
+        "from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, f1_score\n",
+        "from transformers import AutoModelForSequenceClassification, AutoTokenizer, DataCollatorWithPadding, Trainer, TrainingArguments\n",
+        "\n",
+        "PROJECT_ROOT = Path.cwd().resolve().parent if Path.cwd().name == 'notebooks' else Path.cwd().resolve()\n",
+        "MODEL_DIR = PROJECT_ROOT / 'src' / 'models' / 'saved_emotion_model'\n",
+        "REPORT_DIR = PROJECT_ROOT / 'reports' / 'module_2_emotion_classification'\n",
+        "MODEL_NAME = 'distilbert-base-uncased'\n",
+        "\n",
+        "MODEL_DIR.mkdir(parents=True, exist_ok=True)\n",
+        "REPORT_DIR.mkdir(parents=True, exist_ok=True)\n",
+        "sys.path.append(str(PROJECT_ROOT / 'src' / 'models'))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "f5675968",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "dataset = load_dataset('dair-ai/emotion', 'split')\n",
+        "label_names = dataset['train'].features['label'].names\n",
+        "id2label = {i: label for i, label in enumerate(label_names)}\n",
+        "label2id = {label: i for i, label in id2label.items()}\n",
+        "\n",
+        "print(dataset)\n",
+        "print(id2label)\n",
+        "pd.DataFrame(dataset['train'][:5])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "d92e7b96",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
+        "\n",
+        "def tokenize(batch):\n",
+        "    return tokenizer(batch['text'], truncation=True, max_length=128)\n",
+        "\n",
+        "encoded = dataset.map(tokenize, batched=True)\n",
+        "encoded = encoded.rename_column('label', 'labels')\n",
+        "encoded.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])\n",
+        "data_collator = DataCollatorWithPadding(tokenizer=tokenizer)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "de89cb7b",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "model = AutoModelForSequenceClassification.from_pretrained(\n",
+        "    MODEL_NAME,\n",
+        "    num_labels=len(label_names),\n",
+        "    id2label=id2label,\n",
+        "    label2id=label2id,\n",
+        ")\n",
+        "\n",
+        "def compute_metrics(eval_pred):\n",
+        "    logits, labels = eval_pred\n",
+        "    predictions = np.argmax(logits, axis=-1)\n",
+        "    return {\n",
+        "        'accuracy': accuracy_score(labels, predictions),\n",
+        "        'macro_f1': f1_score(labels, predictions, average='macro'),\n",
+        "    }\n",
+        "\n",
+        "training_kwargs = {\n",
+        "    'output_dir': str(PROJECT_ROOT / 'checkpoints' / 'emotion_distilbert'),\n",
+        "    'learning_rate': 2e-5,\n",
+        "    'per_device_train_batch_size': 16,\n",
+        "    'per_device_eval_batch_size': 32,\n",
+        "    'num_train_epochs': 3,\n",
+        "    'weight_decay': 0.01,\n",
+        "    'save_strategy': 'epoch',\n",
+        "    'load_best_model_at_end': True,\n",
+        "    'metric_for_best_model': 'macro_f1',\n",
+        "    'greater_is_better': True,\n",
+        "    'logging_steps': 50,\n",
+        "    'report_to': 'none',\n",
+        "}\n",
+        "\n",
+        "strategy_name = 'eval_strategy' if 'eval_strategy' in signature(TrainingArguments).parameters else 'evaluation_strategy'\n",
+        "training_kwargs[strategy_name] = 'epoch'\n",
+        "\n",
+        "args = TrainingArguments(**training_kwargs)\n",
+        "\n",
+        "trainer_kwargs = {\n",
+        "    'model': model,\n",
+        "    'args': args,\n",
+        "    'train_dataset': encoded['train'],\n",
+        "    'eval_dataset': encoded['validation'],\n",
+        "    'data_collator': data_collator,\n",
+        "    'compute_metrics': compute_metrics,\n",
+        "}\n",
+        "\n",
+        "tokenizer_arg = 'processing_class' if 'processing_class' in signature(Trainer).parameters else 'tokenizer'\n",
+        "trainer_kwargs[tokenizer_arg] = tokenizer\n",
+        "\n",
+        "trainer = Trainer(**trainer_kwargs)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "52e9a3be",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "trainer.train()\n",
+        "validation_metrics = trainer.evaluate(encoded['validation'])\n",
+        "test_output = trainer.predict(encoded['test'])\n",
+        "\n",
+        "test_predictions = np.argmax(test_output.predictions, axis=-1)\n",
+        "test_labels = test_output.label_ids\n",
+        "test_metrics = compute_metrics((test_output.predictions, test_labels))\n",
+        "\n",
+        "print(validation_metrics)\n",
+        "print(test_metrics)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "7c79dfb5",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "model.save_pretrained(MODEL_DIR)\n",
+        "tokenizer.save_pretrained(MODEL_DIR)\n",
+        "\n",
+        "report_text = classification_report(test_labels, test_predictions, target_names=label_names, zero_division=0)\n",
+        "report_dict = classification_report(test_labels, test_predictions, target_names=label_names, output_dict=True, zero_division=0)\n",
+        "matrix = confusion_matrix(test_labels, test_predictions)\n",
+        "\n",
+        "(REPORT_DIR / 'test_classification_report.txt').write_text(report_text, encoding='utf-8')\n",
+        "pd.DataFrame(report_dict).transpose().to_csv(REPORT_DIR / 'test_classification_report.csv')\n",
+        "pd.DataFrame(matrix, index=label_names, columns=label_names).to_csv(REPORT_DIR / 'test_confusion_matrix.csv')\n",
+        "\n",
+        "summary = {\n",
+        "    'base_model': MODEL_NAME,\n",
+        "    'dataset': 'dair-ai/emotion',\n",
+        "    'labels': label_names,\n",
+        "    'validation_metrics': validation_metrics,\n",
+        "    'test_metrics': test_metrics,\n",
+        "    'explainability_method': 'word occlusion on final predicted confidence',\n",
+        "}\n",
+        "(REPORT_DIR / 'metrics_summary.json').write_text(json.dumps(summary, indent=2), encoding='utf-8')\n",
+        "\n",
+        "print(report_text)\n",
+        "print(f'Saved model to: {MODEL_DIR}')\n",
+        "print(f'Saved reports to: {REPORT_DIR}')"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "349facfc",
+      "metadata": {},
+      "source": [
+        "## Explainability Check\n",
+        "\n",
+        "This is not a formal clinical explanation. It is a practical debugging tool: remove one word at a time and measure how much the model's confidence in the predicted emotion drops."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "225e1691",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from emotion_classifier import EmotionClassifier\n",
+        "\n",
+        "classifier = EmotionClassifier(model_dir=MODEL_DIR)\n",
+        "examples = [\n",
+        "    'I feel anxious and overwhelmed and I cannot sleep.',\n",
+        "    'I finally feel hopeful and proud of myself.',\n",
+        "    'I am angry because nobody listens to me.',\n",
+        "]\n",
+        "\n",
+        "explanations = {text: classifier.explain(text, top_k=6) for text in examples}\n",
+        "(REPORT_DIR / 'explanation_examples.json').write_text(json.dumps(explanations, indent=2), encoding='utf-8')\n",
+        "explanations"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python",
+      "pygments_lexer": "ipython3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 5
+}

reports/module_2_emotion_classification/README.md

@@ -0,0 +1,11 @@
+# Module 2 Reports
+
+This folder is populated after running `notebooks/module_2_emotion_training.ipynb` in Colab.
+
+Expected generated files:
+
+- `metrics_summary.json`
+- `test_classification_report.txt`
+- `test_classification_report.csv`
+- `test_confusion_matrix.csv`
+- `explanation_examples.json`

requirements.txt

@@ -3,4 +3,8 @@ scikit-learn==1.9.0
 joblib==1.5.3
 gradio==6.18.0
 datasets==5.0.0
+transformers
+torch
+accelerate
+evaluate
 jupyter

src/models/emotion_classifier.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import argparse
+import json
+import re
+from pathlib import Path
+from typing import Any
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[2]
+DEFAULT_MODEL_DIR = PROJECT_ROOT / "src" / "models" / "saved_emotion_model"
+
+
+def _load_transformer_stack() -> tuple[Any, Any, Any]:
+    try:
+        import torch
+        from transformers import AutoModelForSequenceClassification, AutoTokenizer
+    except ImportError as exc:
+        raise ImportError(
+            "Module 2 requires torch and transformers. Install them with "
+            "`python -m pip install -r requirements.txt`, or run the Colab notebook."
+        ) from exc
+
+    return torch, AutoModelForSequenceClassification, AutoTokenizer
+
+
+class EmotionClassifier:
+    """Transformer emotion classifier with confidence and simple word-occlusion explanations."""
+
+    def __init__(self, model_dir: str | Path = DEFAULT_MODEL_DIR) -> None:
+        self.model_dir = Path(model_dir)
+        self.torch = None
+        self.tokenizer = None
+        self.model = None
+        self.id2label: dict[int, str] = {}
+
+    def load_model(self) -> None:
+        if not self.model_dir.exists():
+            raise FileNotFoundError(
+                f"Emotion model not found at {self.model_dir}. "
+                "Train it first with notebooks/module_2_emotion_training.ipynb."
+            )
+
+        torch, model_cls, tokenizer_cls = _load_transformer_stack()
+        self.torch = torch
+        self.tokenizer = tokenizer_cls.from_pretrained(self.model_dir)
+        self.model = model_cls.from_pretrained(self.model_dir)
+        self.model.eval()
+
+        config_labels = self.model.config.id2label
+        self.id2label = {int(key): value for key, value in config_labels.items()}
+
+    def predict_with_confidence(self, text: str) -> dict[str, Any]:
+        clean_text = text.strip()
+        if not clean_text:
+            return {
+                "emotion": "unknown",
+                "confidence": 0.0,
+                "is_confident": False,
+                "message": "Please enter text to classify.",
+            }
+
+        if self.model is None or self.tokenizer is None or self.torch is None:
+            self.load_model()
+
+        inputs = self.tokenizer(
+            clean_text,
+            return_tensors="pt",
+            truncation=True,
+            max_length=128,
+        )
+
+        with self.torch.no_grad():
+            logits = self.model(**inputs).logits
+            probabilities = self.torch.softmax(logits, dim=-1)[0]
+
+        best_index = int(probabilities.argmax().item())
+        confidence = float(probabilities[best_index].item())
+
+        return {
+            "emotion": self.id2label.get(best_index, str(best_index)),
+            "confidence": confidence,
+            "is_confident": confidence >= 0.60,
+            "message": None,
+        }
+
+    def explain(self, text: str, top_k: int = 8) -> dict[str, Any]:
+        """Estimate influential words by measuring confidence drop after removing each word."""
+        base_prediction = self.predict_with_confidence(text)
+        target_emotion = base_prediction["emotion"]
+        base_confidence = base_prediction["confidence"]
+        words = re.findall(r"\b[\w']+\b", text)
+
+        impacts = []
+        for index, word in enumerate(words):
+            reduced_words = words[:index] + words[index + 1 :]
+            if not reduced_words:
+                continue
+
+            reduced_text = " ".join(reduced_words)
+            reduced_prediction = self.predict_with_confidence(reduced_text)
+            confidence_drop = base_confidence - (
+                reduced_prediction["confidence"]
+                if reduced_prediction["emotion"] == target_emotion
+                else 0.0
+            )
+
+            impacts.append(
+                {
+                    "word": word,
+                    "impact": round(float(confidence_drop), 4),
+                }
+            )
+
+        impacts = sorted(impacts, key=lambda item: item["impact"], reverse=True)
+        return {
+            "prediction": base_prediction,
+            "top_evidence": impacts[:top_k],
+            "method": "word occlusion: larger impact means removing the word reduced confidence more",
+        }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run Module 2 emotion inference.")
+    parser.add_argument("text", nargs="?", default="I feel anxious and overwhelmed today.")
+    parser.add_argument("--explain", action="store_true")
+    parser.add_argument("--model-dir", default=DEFAULT_MODEL_DIR, type=Path)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    classifier = EmotionClassifier(model_dir=args.model_dir)
+    output = classifier.explain(args.text) if args.explain else classifier.predict_with_confidence(args.text)
+    print(json.dumps(output, indent=2))

src/models/emotion_detector_ui.py

@@ -0,0 +1,28 @@
+import gradio as gr
+
+from emotion_classifier import EmotionClassifier
+
+
+classifier = EmotionClassifier()
+
+
+def predict_emotion(text: str) -> dict:
+    return classifier.explain(text or "", top_k=6)
+
+
+interface = gr.Interface(
+    fn=predict_emotion,
+    inputs=gr.Textbox(
+        lines=5,
+        placeholder="Type an English mental-health related message...",
+        label="User Message",
+    ),
+    outputs=gr.JSON(label="Emotion Result"),
+    title="Module 2: Emotion Classification",
+    description="DistilBERT emotion classifier with confidence and word-occlusion explanation.",
+    flagging_mode="never",
+)
+
+
+if __name__ == "__main__":
+    interface.launch()