kiruthik2409/mental-health-xgboost-onnx

Mental Health Prediction Model

This project generates a synthetic dataset, trains a multi-label XGBoost classifier, and exports it to ONNX format for use in web applications.

Files

• train_and_export.py: Main script to generate data, train model, and export ONNX.
• predict.py: Script to manually test the model with your inputs.
• upload_to_hf.py: Script to upload the model to Hugging Face.
• requirements.txt: Python dependencies.
• mental_health_model.onnx: The trained model optimized for inference.
• model_features.json: List of input features expected by the model in order.
• MODEL_ANALYSIS.md: Detailed performance report.

Setup & Run

1. Install Dependencies:

   pip install -r requirements.txt
   

2. Run Training:

   python train_and_export.py
   

   This will:

   • Generate 5,000 synthetic patient records.
   • Train the XGBoost model (predicting Depression, Anxiety, Bipolar, ADHD, SubstanceRelated).
   • Evaluate performance and print metrics.
   • Save mental_health_model.onnx and model_features.json.

output Format

The model predicts probabilities for 5 conditions in this order:

1. Depression
2. Anxiety
3. Bipolar
4. ADHD
5. SubstanceRelated

Using in React (JavaScript)

To use this model in your React application, use onnxruntime-web.

1. Install ONNX Runtime

npm install onnxruntime-web

2. Copy Files

Copy mental_health_model.onnx to your public folder (e.g., public/models/mental_health_model.onnx).

3. Inference Code Example

import * as ort from 'onnxruntime-web';

// Feature list must match model_features.json order
const FEATURES = [
  'Age', 'Sex', 
  'insomnia', 'fatigue', 'panic', 'panic_sleep', 'social_withdrawal',
  'mood_swings', 'anhedonia', 'rumination', 'appetite_change', 'concentration_issues',
  'self_harm_ideation', 'irritability', 'withdrawal', 'anxiety_general', 'psychomotor',
  'suicidal_thoughts', 'hopelessness', 'panic_physical'
];

async function predictMentalHealth(inputs) {
  try {
    // 1. Create session
    const session = await ort.InferenceSession.create('/models/mental_health_model.onnx');

    // 2. Prepare input tensor
    // Inputs must be float32
    const inputData = Float32Array.from(FEATURES.map(f => inputs[f] || 0));
    
    // Shape: [1, 20] (Batch size 1, 20 features)
    const tensor = new ort.Tensor('float32', inputData, [1, 20]);

    // 3. Run inference
    // Feeds: object key 'float_input' must match the name defined in python script
    const feeds = { float_input: tensor };
    const results = await session.run(feeds);

    // 4. Parse results
    // Output name depends on sklearn conversion, usually 'probabilities' or 'output_probability'
    // For MultiOutput, it might return multiple outputs or a combined sequence.
    // Check results object keys. Usually: results.label, results.probabilities
    
    console.log("Prediction Results:", results);
    
    // Note: The structure of 'results' varies by ONNX converter version. 
    // You might get a map or a list of tensors. Inspect 'results' in console.
    
    return results;

  } catch (e) {
    console.error("Inference failed:", e);
  }
}

Note: Dealing with MultiOutputClassifier in ONNX can produce complex output structures (e.g., a sequence of maps). You may need to inspect session.run() output to map probabilities back to the 5 conditions.

Hugging Face Integration

Uploading to Hugging Face

1. Log in to your Hugging Face account via CLI:

   huggingface-cli login
   

2. Run the upload script:

   python upload_to_hf.py
   

   Enter your username when prompted. The script will create a repository mental-health-xgboost-onnx and upload the artifacts.

Downloading for Use in Project

You can download the files directly or use huggingface_hub in Python.

Method 1: Direct Download (Web App)

• Go to your model page: https://huggingface.co/<your-username>/mental-health-xgboost-onnx
• Go to "Files and versions" tab.
• Download mental_health_model.onnx and place it in your project's public folder.

Method 2: Python

from huggingface_hub import hf_hub_download

model_path = hf_hub_download(repo_id="<your-username>/mental-health-xgboost-onnx", filename="mental_health_model.onnx")
print(f"Model downloaded to: {model_path}")