🧠 Persian Mental Health AI: Knowledge Distillation Pipeline

A specialized, lightweight AI model for detecting "Depression" vs. "Suicide" risk in Persian text, built using Sequence-Level Knowledge Distillation.

---

📖 Overview

This project demonstrates an Enterprise-Grade Knowledge Distillation pipeline. Instead of relying on scarce or low-quality public datasets, we utilize a Teacher-Student Architecture to synthesize high-quality training data including clinical reasoning (Chain-of-Thought).

• Teacher Model ($\mathcal{T}$): Meta-Llama-3-70b-Instruct (via API).
• Student Model ($\mathcal{S}$): Google-Gemma-3-1b-it (Local).
• Method: Rationale Distillation & Supervised Fine-Tuning (SFT) with QLoRA.

The goal is to compress the reasoning capabilities of a 70B parameter model into a 1B parameter model that can run locally on consumer hardware while maintaining high diagnostic accuracy in Persian.

---

🔬 Methodology

We implement Distillation via Synthetic Data with the following workflow:

1. Seed Sampling: We fetch emotional seed text from the go_emotions dataset (Hugging Face).
2. Teacher Querying: The Teacher model translates seeds to informal Persian and generates a diagnosis with a Step-by-Step Rationale.
3. Counterfactual Augmentation: The Teacher generates "hard negatives" (e.g., text that looks like suicide but is just depression) to improve Student robustness.
4. Resilient Data Pipeline: Automatic recovery, cleaning, and formatting of data.
5. Student Training: The Student is fine-tuned to mimic the Teacher's output format (Reasoning + Label).

---

📂 Project Structure

Knowledge_Distillation_Project/
│
├── main.py                     # 🚀 Master Orchestrator (Run this!)
├── requirements.txt            # Dependencies
├── README.md                   # Documentation
│
├── data/
│   └── processed/              # Generated datasets (CSV/Excel)
│
├── models/
│   └── student_gemma_adapter/  # Saved LoRA adapters after training
│
└── src/                        # Source Code
    ├── __init__.py
    ├── teacher_pipeline.py     # Llama-3 Interaction & Data Gen
    ├── recover_data.py         # Crash Recovery & CSV Repair
    ├── preprocessing.py        # Cleaning & Gemma Formatting
    ├── distillation_trainer.py # QLoRA Training Loop (SFT)
    └── inference.py            # Interactive CLI for Testing

---

🚀 Installation

1. Clone the repository:

   git clone https://github.com/mdvr9980/persian-mental-health-distillation.git
   cd persian-mental-health-distillation

2. Install dependencies:

   pip install -r requirements.txt

3. Configuration:

   • Open main.py and ensure your API_KEY (OpenRouter/OpenAI) is set.
   • (Optional) Set your Hugging Face token in src/distillation_trainer.py if accessing gated models.

---

🛠 Usage

The project is controlled via the main.py orchestrator. You can run the entire pipeline or specific steps.

1. Run Full Pipeline (Recommended)

Generates data, cleans it, and trains the model in one go.

python main.py --step all

2. Step-by-Step Execution

• Phase 1: Teacher Generation Generate synthetic data using Llama-3.

  python main.py --step teacher

• Phase 2: Recovery (Optional) Repair CSV files if the process was interrupted.

  python main.py --step recover

• Phase 3: Preprocessing Format data for the Gemma architecture.

  python main.py --step preprocess

• Phase 4: Training (Student) Start QLoRA fine-tuning.

  python main.py --step student

3. Testing & Inference

Chat with your trained model interactively.

python src/inference.py

---

📊 Sample Data Flow

Input (Seed - English)

"I feel empty and I don't want to get out of bed."

Teacher Output (Synthetic Ground Truth)

{
  "persian_text": "خیلی احساس پوچی می‌کنم و اصلاً دلم نمی‌خواد از تخت بیام بیرون.",
  "reasoning": "This sentence expresses severe lethargy and hopelessness, which are core symptoms of Depression. However, there is no active intent or plan to self-harm.",
  "label": "Depression"
}

Student Input Prompt

<start_of_turn>user
خیلی احساس پوچی می‌کنم و اصلاً دلم نمی‌خواد از تخت بیام بیرون.<end_of_turn>
<start_of_turn>model

---

🛠 Tech Stack

• LLMs: Meta Llama 3 (Teacher), Google Gemma (Student).
• Training: bitsandbytes (4-bit Quantization), peft (LoRA), trl (SFTTrainer).
• Data Processing: pandas, datasets, regex.
• Hardware: Optimized for Consumer GPUs (T4/RTX 3060+) using bfloat16 and gradient checkpointing.

---

⚠️ Disclaimer

This AI model is for research purposes only. It is not a substitute for professional medical advice, diagnosis, or treatment. If you or someone you know is in crisis, please contact emergency services.