Since Russia launched a full-scale military invasion into Ukraine on February 24, 2022, fighting has caused over one hundred civilian casualties and pushed tens of thousands of Ukrainians to flee to neighboring countries—including Poland. During the war, people can be exposed to many different traumatic events. That raises the chances of developing mental health problems—like post-traumatic stress disorder (PTSD), anxiety, and depression—and poorer life outcomes as adults. Sometimes all they need is someone to talk to, someone to make them feel better, and someone to share their feelings.

What it does

We fine-tuned a 90 Million parameter language model on a custom dataset using AWS EC2 DL1.24xlarge instance powered by 8 Habana Gaudi accelerators. We benchmarked the training with Google Colab Pro which uses Tesla P100 which has 16GB of HBM2 VRAM. The 8 Gaudi accelerator outperformed on the training time and efficiency.

How we built it

We built it by months of research on large language model fine-tuning on accelerators. First made a custom script that helps us to train any custom dataset on any task for our language model, then we trained on Google Colab Pro which uses Tesla P100. After we noted the performance, trained the same on DL1.24xlarge and benchmarked both of these including the usage and consistency in conversations. We introduced Data-Parallel in Inorder to train on multiple accelerators and conducted multiple brainstorming sessions with people from the All India Institute Of Medical Science for understanding the mental health problems faced by people during the war.

Custom Dataset Sample

text:I am feeling really sad?   labels:why are you feeling sad?
text:A war is going on here in Ukrain?   labels:yes, I know. Russia and Ukraine have had no formal diplomatic relations since 24 February 2022. The Russian Federation and Ukraine are currently in a state of war: the Russo-Ukrainian War began in 2014 following the Russian annexation of Crimea from Ukraine. In February 2022, Russia invaded Ukraine across a broad front.
text:Indeed. When will it end  labels:Hope it end soon
text:I am feeling really scared   labels:Don't be scared. Be strong and be brave
text:How can I feel secured, feeling so insecured now   labels:I am here for you. Don't feel scared. We will get through it.
text:I am in a bunker now   labels:That is scary. But things will be better soon. Don't lose hope.
text:thank you   labels:you are welcome. Hope everything works fine.
text:bye   labels:bye.   episode_done=True

Challenges we ran into

Installing the dependencies on the instance has taken a couple of days. Coding the custom script has taken weeks to complete.

Accomplishments that we're proud of

Built the fastest trainable most powerful language model on Gaudi accelerator which will help the humanity

What we learned

Training on Gaudi accelerators. Building a complex language model training pipeline

What's next for Large Language Training For Mental Health

The trained language model can be deployed as an API. This means we can even build a mental health helpline number with our trained conversational AI as a backend to help them feel connected.

Built With

Share this project: