This article was produced as part of the final project for Harvard’s AC215 Fall 2024 course.
Project Github Repo — https://github.com/jxu0510/ac215_closedai
Video — <Link>
- Introduction
- Data Pipeline
- Model Training
- API and Frontend
- Deployment and Scaling
- Future work
- References
Rates of depression and anxiety among young people, particularly college students, have been trending upward over the past decade. With the increasing demand for mental health services, the availability of professional care remains constrained by factors like limited appointment slots and long wait times, often leaving individuals waiting for help. Recognizing this gap, our team developed EmpathOS, a mental health chatbot designed to provide immediate, confidential support and connect users to professional resources when necessary.
In this blog post, we’ll dive into the technical approach and the impact of our project.
Datasets
For our mental health chatbot, we utilized two primary datasets tailored to different stages of the project:
- Mental Health Q&A Dataset
This dataset comprises 658 mental health-related Q&A conversations, focusing on topics such as anxiety, depression, and general therapeutic advice. It consists of 594 training samples and 64 test samples. The data was collected from publicly available FAQs, classical therapy conversations, and general advice discussions, structured to enhance the chatbot’s ability to provide empathetic and contextually relevant responses. - Academic Papers on Mental Health Counseling
This dataset contains text files of academic papers targeting teenagers and young adults, sourced to enhance the chatbot’s ability to generate evidence-based responses. The dataset focuses on topics like stress management, anxiety interventions, and counseling methods.
Data Preprocessing
We developed a containerized approach to handle various types of datasets, each with its own tailored workflow. For Q&A datasets, we transformed raw conversational data into well-structured JSONL files, ready for training and evaluation. Meanwhile, for our Retrieval-Augmented Generation (RAG) process, we used raw text documents and turned them into embedded chunks stored in a ChromaDB vector database. By separating and streamlining these preprocessing steps, we can maintain data security and integrate multiple data sources into our model pipeline.
Data Versioning
We adopted Data Version Control (DVC) to manage dataset iterations and ensure reproducibility.
- Configuration: Separate remotes for the RAG dataset.
- Versioning: Tracks changes to the raw datasets.
- Storage: Versioned datasets and metadata are stored on GCS and locally.
This system allows us to compare iterations, revert to previous versions, and ensure collaboration transparency.
Model Summary
For this project, we fine-tuned the GEMINI model and utilized a Retrieval-Augmented Generation (RAG) workflow to build a robust mental health chatbot. GEMINI is a large language model specifically designed for conversational AI, offering a strong base for generating empathetic and contextually relevant responses. RAG, on the other hand, enhances the chatbot by combining generative capabilities with external knowledge retrieval, which ensures accurate and evidence-based responses.
The GEMINI model was fine-tuned on our mental health Q&A dataset, and for the RAG workflow, we fed the system academic papers specifically focused on mental health counseling for teenagers and young adults.
Training Pipeline
We implemented a containerized training pipeline to streamline the fine-tuning process. The Model Fine-Tuning Container preprocesses the mental health dataset, executes the fine-tuning task, and deploys the model to a Google Cloud endpoint.
For the RAG workflow, we prepared data by chunking academic papers into smaller text blocks, embedding them using a vector database (ChromaDB), and integrating the database with the fine-tuned model. This setup ensures that the chatbot retrieves and generates responses supported by psychological research.
Application Design
We developed design documents that outline the high-level architecture of our mental health chatbot application. These include a solution architecture, which presents the overall strategy of the project, and a technical architecture, which details the interactions between various components.
The solution architecture in Figure 1 illustrates the flow of processes (tasks handled by developers and users), execution (operations across the project pipeline), and state (data and artifacts stored at different stages). This view abstracts technical details to provide a high-level understanding of the system.
The technical architecture in Figure 2 provides implementation details, showing components responsible for specific actions and the communication between them.
Frontend
Our frontend is designed to provide a clean, intuitive interface that feels effortless to use. It uses a simple and familiar chat-style interface where users can input their thoughts, questions, or feelings and then submit them to our backend system. Its minimalist design reduces visual distractions, allowing users to connect with the system and get the support or information they need without unnecessary complexity.
Backend
We used Flask to build our backend endpoints, facilitating interaction between the frontend interface and our underlying data and model infrastructure. We implemented the following endpoints:
- /: A GET method that renders the main index.html page, providing the user with our chat interface.
- /chat: A POST method that accepts a user’s message as JSON input. The backend uses Vertex AI embeddings to translate the user’s query into a high-dimensional vector, queries a ChromaDB vector database for relevant document chunks, and then passes the combined information to a Vertex AI generative model to produce a contextually grounded response.
We used Ansible to automate deployment and leveraged Kubernetes for scaling. By containerizing our application and using Ansible playbooks, we can easily provision infrastructure, configure virtual machines, and set up Kubernetes clusters that automatically balance workloads. This approach supports both traditional VM-based deployments and container orchestration via Kubernetes, enabling quick scalability as traffic grows. The use of a standardized deployment container simplifies tool management and ensures that our development, staging, and production environments remain consistent and stable throughout the continuous integration and delivery cycle.
- Model Optimization and Continuous Improvement: We will continue refining our language models and embedding methods, integrating feedback to enhance accuracy, relevance, and efficiency over time.
- Conduct UX Interviews and Gather Feedback: We will interview selected early trial users to identify usability issues, improve the interface, and explore new features that align with real-world needs.
- Expand Data Sources and Domain Coverage: We intend to broaden our dataset for RAG by incorporating more specialized content and considering multilingual capabilities to address a wider range of user queries.
- Atkins, M. S., Hoagwood, K. E., Kutash, K., & Seidman, E. (2010). Toward the integration of education and mental health in schools. Administration and Policy in Mental Health and Mental Health Services Research, 37(1–2), 40–47. https://doi.org/10.1007/s10488-010-0299-7
- Cavioni, V., Grazzani, I., Ornaghi, V., Agliati, A., & Pepe, A. (2021). Adolescents’ mental health at school: The mediating role of life satisfaction. Frontiers in Psychology, 12, 720628. https://doi.org/10.3389/fpsyg.2021.720628
- Dowdy, E., Ritchey, K., & Kamphaus, R. W. (2010). School-based screening: A population-based approach to inform and monitor children’s mental health needs. School Mental Health, 2(4), 166–176. https://doi.org/10.1007/s12310-010-9036-3
- Fazel, M., Hoagwood, K., Stephan, S., & Ford, T. (2014). Mental health interventions in schools in high-income countries. The Lancet Psychiatry, 1(5), 377–387. https://doi.org/10.1016/S2215-0366(14)70312-8
- GarcÃa-Carrión, R., Villarejo-Carballido, B., & Villardón-Gallego, L. (2019). Children and adolescents mental health: A systematic review of interaction-based interventions in schools and communities. Frontiers in Psychology, 10, 918. https://doi.org/10.3389/fpsyg.2019.00918
- Nakao, M., Shirotsuki, K., & Sugaya, N. (2021). Cognitive-behavioral therapy for management of mental health and stress-related disorders: Recent advances in techniques and technologies. BioPsychoSocial Medicine, 15, 2. https://doi.org/10.1186/s13030-021-00219-w