Update README.md

README.md

@@ -17,6 +17,9 @@ tags:
 - emotionalsupport
 - mentalsupport
 - advisor
+- medical
+- not-for-all-audiences
+pipeline_tag: text-generation
 ---
 
 # Model Card for Model ID
@@ -25,7 +28,8 @@ tags:
 
 Falcon-1B-Mental-Health-Advisor 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.
 
-## Model Details
+# 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)
 
@@ -35,27 +39,151 @@ This is a LoRA adapter for the Falcon-RW-1B model. It was fine-tuned on the 'mar
 
 Since this model is an adapter, it **must** be loaded with the original Falcon-RW-1B model using PEFT:
 
+### Dependencies
+```bash
+pip install transformers accelerate torch peft bitsandbytes language_tool_python
+```
+
+### Basic Usage
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
 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"
-model = AutoModelForCausalLM.from_pretrained(base_model, device_map="auto")
-model = PeftModel.from_pretrained(model, "ShivomH/Falcon-1B-Finetuned-Mental-Health")
-
-tokenizer = AutoTokenizer.from_pretrained("ShivomH/Falcon-1B-Finetuned-Mental-Health")
+peft_model = "ShivomH/Falcon-1B-Finetuned-Mental-Health"
+
+# 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.3,
+                top_p=0.9,
+                repetition_penalty=1.4,
+                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.
+
 ### Model Sources [optional]
 
 <!-- Provide the basic links for the model. -->
@@ -64,6 +192,7 @@ tokenizer = AutoTokenizer.from_pretrained("ShivomH/Falcon-1B-Finetuned-Mental-He
 - **Paper [optional]:** [More Information Needed]
 - **Demo [optional]:** [More Information Needed]
 
+### Framework versions
 
 - PEFT 0.14.0
 ```