Falcon-1B-Mental-Health-v2

Model Card for Model ID

<!-- Provide a quick summary of what the model is/does. -->

This is the version-2 of Falcon-1B-Mental-Health. Falcon-1B-Mental-Health-v2 is a fine-tuned version of the tiiuae/falcon-rw-1b model, adapted for providing empathetic and contextually relevant responses to mental health-related queries. The model has been trained on a curated dataset to assist in mental health conversations, offering advice, guidance, and support for individuals dealing with issues like stress, anxiety, and depression. It provides a compassionate approach to mental health queries while focusing on promoting emotional well-being and mental health awareness.

Important Note

Mental Health is a sensitive topic. Preferably, use the code snippet provided below in order to get optimal results.

Falcon-1B Fine-Tuned for Mental Health (LoRA)

This is the version-2 of Falcon-1B-Mental-Health. In this version, the primary attempt was to improve the performance of the model by making slight modifications in the LoRA config and training arguements. This is a LoRA adapter for the Falcon-RW-1B model. It was fine-tuned on the 'marmikpandya/mental-health' dataset.

Usage

Since this model is an adapter, it must be loaded with the original Falcon-RW-1B model using PEFT:

Dependencies

pip install transformers accelerate torch peft bitsandbytes language_tool_python

Basic Usage

from peft import PeftModel
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
import torch
import re
import language_tool_python

base_model = "tiiuae/falcon-rw-1b"
peft_model = "ShivomH/Falcon-1B-Mental-Health-v2"

# Load the base model (without LoRA weights initially)
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    torch_dtype=torch.float16,
    device_map="auto"
)

# Load LoRA weights into the model
model = PeftModel.from_pretrained(model, peft_model)

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model)
tokenizer.pad_token = tokenizer.eos_token

## How to Get Started with the Model

# Move the model to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Load the grammar correction tool
tool = language_tool_python.LanguageTool("en-US")
def correct_grammar(text):
    return tool.correct(text)

# --- Safety Filters ---
CRISIS_KEYWORDS = [
    "suicide", "self-harm", "overdose", "addict", "abuse", "rape", "assault", "emergency", "suicidal"
]
CRISIS_RESPONSE = (
    "\n\nIf you're in crisis, please contact a professional immediately. "
    "You can reach the National Suicide Prevention Lifeline at 988 or 112."
    "Please reach out to a trusted friend, family member, or mental health professional. "
    "If you're in immediate danger, consider calling a crisis helpline. Your life matters, and support is available. 🙏"
)

def filter_response(response: str, user_input: str) -> str:
    # Remove URLs, markdown artifacts, and unwanted text
    response = re.sub(r'http\S+', '', response)
    response = re.sub(r'\[\w+\]|\(\w+\)|\*|\#', '', response)
    response = response.split("http")[0].split("©")[0]

    # Enforce brevity: Keep only the first two sentences
    sentences = re.split(r'(?<=[.!?])\s+', response)
    response = " ".join(sentences[:2])  # Keep only first 2 sentences

    # Append crisis response if keywords detected
    if any(keyword in user_input.lower() for keyword in CRISIS_KEYWORDS):
        response += CRISIS_RESPONSE

    # Correct grammar
    response = correct_grammar(response)

    return response

def chat():

    print("Chat with your fine-tuned Falcon model (type 'exit' to quit):")

    system_instruction = (
        "You are an empathetic AI specialized in mental health support. "
        "Provide short, supportive, and comforting responses. "
        "Validate the user's emotions and offer non-judgmental support. "
        "If a crisis situation is detected, suggest reaching out to a mental health professional immediately. "
        "Your responses should be clear, concise, and free from speculation. "
        # "Examples:\n"
        # "User: I feel really anxious lately.\n"
        # "AI: I'm sorry you're feeling this way. Anxiety can be overwhelming, but you're not alone. Would you like to try some grounding techniques together?\n\n"
        # "User: I haven't been able to sleep well.\n"
        # "AI: That sounds frustrating. Sleep troubles can be tough. Have you noticed anything that helps, like adjusting your bedtime routine?\n"
    )

    # Store short chat history for context
    chat_history = []

    while True:
        user_input = input("\nYou: ")
        if user_input.lower() == "exit":
            break

        # Maintain short chat history (last 2 exchanges)
        chat_history.append(f"User: {user_input}")
        chat_history = chat_history[-2:]

        # Structure prompt
        prompt = f"{system_instruction}\n" + "\n".join(chat_history) + "\nAI:"
        inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")

        with torch.no_grad():
            output = model.generate(
                **inputs,
                max_new_tokens=75,
                pad_token_id=tokenizer.eos_token_id,
                temperature=0.4,
                top_p=0.9,
                repetition_penalty=1.2,
                do_sample=True,
                no_repeat_ngram_size=2,
                early_stopping=True
            )

        response = tokenizer.decode(output[0], skip_special_tokens=True).split("AI:")[-1].strip()
        response = filter_response(response, user_input)
        print(f"AI: {response}")

        # Store AI response in history
        chat_history.append(f"AI: {response}")

chat()

Model Description

<!-- Provide a longer summary of what this model is. -->

• Developed by: Shivom Hatalkar
• Model type: Text-generation
• Language(s) (NLP): English
• License: apache-2.0
• Finetuned from model [optional]: falcon-rw-1b

Bias, Risks, and Limitations

• Not a Substitute for Professional Care: This model is not a licensed mental health professional. Its responses may be incomplete, inaccurate, or unsuitable for serious conditions.
• Inherent Biases - May reflect biases in training data (e.g., cultural assumptions, stigmatizing language).
• Crisis Limitations - Not designed for crisis intervention (e.g., suicidal ideation, self-harm). Always direct users to human professionals or emergency services.
• Over-Reliance Risk - Outputs could inadvertently worsen symptoms if users interpret them as definitive advice.
• Intended Use - Assist with general emotional support, not diagnosis or treatment.

Training Hyperparameters

Model Sources [optional]

<!-- Provide the basic links for the model. -->

• Repository: [More Information Needed]
• Paper [optional]: [More Information Needed]
• Demo [optional]: [More Information Needed]

Framework versions

• PEFT 0.14.0