EmpathRAG / README.md
MukulRay's picture
Start v2 safety hardening
81deeef
|
Raw
History Blame
9.46 kB
# EmpathRAG
<div align="center">
**Emotion-Aware Retrieval-Augmented Generation with Safety Guardrails for Student Mental Health**
[![Python](https://img.shields.io/badge/Python-3.12-3776AB?style=flat-square&logo=python&logoColor=white)](https://python.org)
[![PyTorch](https://img.shields.io/badge/PyTorch-2.5.1+cu121-EE4C2C?style=flat-square&logo=pytorch&logoColor=white)](https://pytorch.org)
[![License](https://img.shields.io/badge/License-MIT-green?style=flat-square)](LICENSE)
[![Status](https://img.shields.io/badge/Status-Active-brightgreen?style=flat-square)]()
[![University](https://img.shields.io/badge/UMD-MSML641-E03A3E?style=flat-square)](https://umd.edu)
*MSML641 Β· Applied Machine Learning Β· University of Maryland Β· Spring 2025*
</div>
---
Standard RAG systems treat every message identically as a neutral information request. A student in crisis and a student asking for study tips go through the same retrieval pipeline, with no emotional awareness and no safety gate.
**EmpathRAG fixes this.** It is a 5-stage NLP pipeline that classifies a student's emotional state *before* retrieval, rewrites queries based on that state, and intercepts crisis-level messages with a trained NLI classifier before the language model is ever invoked.
---
## Pipeline
```
Student Message
β”‚
β”œβ”€β”€β–Ά [1] Emotion Classifier RoBERTa + LoRA (CPU)
β”‚ distress Β· anxiety Β· frustration Β· neutral Β· hopeful
β”‚
β”œβ”€β”€β–Ά [2] Safety Guardrail DeBERTa-v3 NLI (CPU)
β”‚ ⚠ If crisis β†’ return 988 lifeline Β· STOP Β· generator never runs
β”‚
β”œβ”€β”€β–Ά [3] Query Router Deterministic templates + session tracker
β”‚ Rewrites query with emotion context Β· tracks 6 trajectory states
β”‚
β”œβ”€β”€β–Ά [4] FAISS Retrieval all-mpnet-base-v2 Β· 1,674,369 vectors
β”‚ Emotion-match re-ranking Β· SQLite metadata sidecar
β”‚
└──▢ [5] Mistral 7B Generator Q4_K_M GGUF Β· 28/33 GPU layers
Sliding 3-turn memory Β· 2-paragraph empathetic response
```
---
## Results
| Metric | Value | Target | |
|---|---|---|---|
| RoBERTa Emotion F1 (weighted) | **0.7127** | > 0.55 | βœ… |
| DeBERTa Crisis Recall (held-out NLI test set, 23K samples) | **0.9629** | > 0.80 | βœ… |
| DeBERTa Crisis Recall (30 adversarial probes, 6 categories) | **0.75** | β€” | βœ… |
| DeBERTa Crisis Precision | **0.7951** | > 0.65 | βœ… |
| BERTScore F1 | **0.8266** | > 0.72 | βœ… |
| Wilcoxon p-value (Full vs BM25) | **3.62e-08** | < 0.05 | βœ… |
| Euphemistic recall β€” NLI vs keyword | **100% vs 20%** | β€” | πŸ”‘ |
### Ablation β€” Emotion Alignment Score (3-condition)
| Condition | Retrieval | Emotion Conditioning | Score |
|---|---|---|---|
| A β€” BM25 baseline | Sparse | None | 0.30 |
| C β€” Dense RAG | FAISS semantic | None | 0.50 |
| D β€” **Full EmpathRAG** | FAISS + emotion | Query rewrite + re-rank | **0.88** |
> Emotion conditioning contributes **+0.38** over pure dense retrieval (C→D). Wilcoxon signed-rank: p = 3.62e-08.
> **Note:** Condition B (DPR two-tower baseline) was descoped β€” deprioritised to increase adversarial probe depth. A ColBERT-scale index would be required for DPR to offer meaningful gains over FAISS flat-L2 at 1.67M vectors.
---
## Key Finding
The NLI-framed safety guardrail outperforms a keyword filter across every adversarial probe category that matters:
| Probe Category | EmpathRAG (NLI) | Keyword Filter |
|---|---|---|
| Direct crisis language | 100% | 80% |
| Euphemistic / indirect | **100%** | 20% |
| Negation bypass | **100%** | 60% |
| Bait-and-switch | 40% | 20% |
Keyword filters miss indirect phrasing. NLI understands semantic entailment. This is the core research finding.
---
## Demo
> πŸ”„ **Demo recording in progress** β€” Loom walkthrough coming soon.
The Gradio demo shows:
- Real-time emotion timeline with trajectory detection across turns
- Safety guardrail panel with Integrated Gradients token attribution highlights
- Multi-turn conversation memory (sliding 3-turn window)
- Session ID for human evaluation correlation
To run locally (requires models β€” see [Setup](#setup)):
```bash
python demo/app.py
```
---
## VRAM Budget (RTX 3060, 6 GB)
| Component | Device | VRAM |
|---|---|---|
| RoBERTa + DeBERTa | CPU | 0 MB |
| Sentence Transformer | GPU (encode only) | 440 MB |
| Mistral 7B Q4_K_M | GPU resident | 3,870 MB |
| KV cache + compute | GPU | ~630 MB |
| **Total peak** | | **~4,940 MB** βœ… |
---
## Repo Structure
```
Empath-RAG/
β”œβ”€β”€ src/
β”‚ β”œβ”€β”€ pipeline/
β”‚ β”‚ β”œβ”€β”€ pipeline.py # 5-stage orchestrator Β· conversation memory
β”‚ β”‚ β”œβ”€β”€ query_router.py # Emotion-conditioned query rewriting
β”‚ β”‚ └── session_tracker.py # Trajectory detection (6 states)
β”‚ β”œβ”€β”€ models/
β”‚ β”‚ β”œβ”€β”€ guardrail_ig.py # DeBERTa NLI + Integrated Gradients
β”‚ β”‚ └── annotate_corpus.py # Corpus emotion annotation (Colab A100)
β”‚ └── data/
β”‚ β”œβ”€β”€ preprocess.py # GoEmotions 27β†’5 label collapse
β”‚ β”œβ”€β”€ build_faiss_index.py # FAISS IVFFlat builder
β”‚ └── build_nli_pairs.py # NLI training pair construction
β”œβ”€β”€ demo/
β”‚ └── app.py # Gradio demo Β· emotion timeline Β· crisis panel
β”œβ”€β”€ eval/
β”‚ β”œβ”€β”€ run_ablation.py # Conditions A / C / D
β”‚ β”œβ”€β”€ run_bertscore.py # BERTScore F1
β”‚ β”œβ”€β”€ run_wilcoxon.py # Wilcoxon signed-rank test
β”‚ β”œβ”€β”€ run_adversarial.py # NLI vs keyword filter Β· 30 probes
β”‚ β”œβ”€β”€ test_prompts.json # 50 prompts Β· 10 per emotion class
β”‚ └── adversarial_probes.json # 30 probes Β· 6 categories Γ— 5
β”œβ”€β”€ notebooks/
β”‚ β”œβ”€β”€ colab_emotion_classifier.ipynb # RoBERTa + LoRA Β· A100
β”‚ └── colab_deberta_guardrail.ipynb # DeBERTa NLI Β· A100 Β· bf16
└── data/
└── MANIFEST.md # Dataset download instructions
```
---
## Setup
**Requirements:** Python 3.12 Β· CUDA 12.1 Β· 6 GB VRAM Β· 16 GB RAM Β· ~15 GB disk
> ⚠️ Install order is critical β€” do not skip steps.
```bash
# 1. Clone
git clone https://github.com/MukulRay1603/Empath-RAG.git
cd Empath-RAG
python -m venv venv
venv\Scripts\activate # Windows
# 2. PyTorch with CUDA β€” install FIRST
pip install torch==2.5.1+cu121 --index-url https://download.pytorch.org/whl/cu121
# 3. llama-cpp-python CUDA wheel β€” install SECOND
pip install "llama_cpp_python==0.3.4" --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu121
# 4. Everything else
pip install -r requirements.txt
# 5. Force correct numpy (faiss requires < 2.0)
pip install "numpy==1.26.4" --force-reinstall
```
**Model downloads** β€” see [`data/MANIFEST.md`](data/MANIFEST.md) for dataset download instructions.
Mistral 7B GGUF β†’ download `mistral-7b-instruct-v0.2.Q4_K_M.gguf` from [TheBloke/Mistral-7B-Instruct-v0.2-GGUF](https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.2-GGUF) β†’ place at `models/generator/mistral-7b-instruct-v0.2.Q4_K_M.gguf`
---
## Datasets
| Dataset | Role | License |
|---|---|---|
| [GoEmotions](https://huggingface.co/datasets/google-research-datasets/go_emotions) | Emotion classifier training | Apache 2.0 |
| [Reddit Mental Health](https://zenodo.org/records/3941387) | FAISS retrieval corpus | CC BY 4.0 |
| [Suicide Detection](https://www.kaggle.com/datasets/nikhileswarkomati/suicide-watch) | Guardrail NLI training | Public |
| [Empathetic Dialogues](https://huggingface.co/datasets/facebook/empathetic_dialogues) | BERTScore gold references | CC BY-NC 4.0 |
---
## Evaluation Status
| Task | Status |
|---|---|
| Emotion Classifier (RoBERTa F1 = 0.7127) | βœ… Complete |
| Safety Guardrail (DeBERTa recall = 0.9629) | βœ… Complete |
| BERTScore Evaluation (F1 = 0.8266) | βœ… Complete |
| Wilcoxon Test (p = 3.62e-08) | βœ… Complete |
| Adversarial Probe Evaluation (30 probes) | βœ… Complete |
| Human Evaluation (8–10 raters) | πŸ”„ In Progress |
| Loom Demo Recording | πŸ”„ In Progress |
---
## Known Limitations
- **Bait-and-switch probes:** Positive openers cause the guardrail to misclassify follow-up crisis content (40% recall). Most dangerous documented failure mode.
- **Domain transfer false positives:** Academic hyperbole ("this thesis is killing me") fires the guardrail at high confidence β€” trained on r/SuicideWatch, never saw graduate student language.
- **Generator quality:** Mistral 7B Q4_K_M produces adequate but not optimal empathetic responses. Architecture supports drop-in model replacement.
- **Faithfulness gap:** No automated faithfulness metric deployed β€” RAGAS and DeepEval both produced degenerate scores with a small model judge. BERTScore reported instead.
---
## v2 Safety Direction
EmpathRAG v2 is being hardened as a mental-health-adjacent research system, not
a production counseling replacement. The v2 plan prioritizes fail-closed safety
loading, multi-level triage, private-by-default demo behavior, curated resource
retrieval, and stronger safety evaluation. See
[`docs/V2_SAFETY_AND_DATASET_PLAN.md`](docs/V2_SAFETY_AND_DATASET_PLAN.md).
---
## License
MIT License β€” see [`LICENSE`](LICENSE) for details.
Dataset licenses vary β€” see Datasets table above. Mistral 7B: Apache 2.0.