Views
No views yet
pip install transformers accelerate torch peft bitsandbytes language_tool_python1from peft import PeftModel
2from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
3import torch
4import re
5import language_tool_python
6
7base_model = "tiiuae/falcon-rw-1b"
8peft_model = "ShivomH/Falcon-1B-Mental-Health-v2"
9
10# Load the base model (without LoRA weights initially)
11model = AutoModelForCausalLM.from_pretrained(
12 base_model,
13 torch_dtype=torch.float16,
14 device_map="auto"
15)
16
17# Load LoRA weights into the model
18model = PeftModel.from_pretrained(model, peft_model)
19
20# Load the tokenizer
21tokenizer = AutoTokenizer.from_pretrained(base_model)
22tokenizer.pad_token = tokenizer.eos_token
23
24## How to Get Started with the Model
25
26# Move the model to GPU if available
27device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
28model.to(device)
29
30# Load the grammar correction tool
31tool = language_tool_python.LanguageTool("en-US")
32def correct_grammar(text):
33 return tool.correct(text)
34
35# --- Safety Filters ---
36CRISIS_KEYWORDS = [
37 "suicide", "self-harm", "overdose", "addict", "abuse", "rape", "assault", "emergency", "suicidal"
38]
39CRISIS_RESPONSE = (
40 "\n\nIf you're in crisis, please contact a professional immediately. "
41 "You can reach the National Suicide Prevention Lifeline at 988 or 112."
42 "Please reach out to a trusted friend, family member, or mental health professional. "
43 "If you're in immediate danger, consider calling a crisis helpline. Your life matters, and support is available. 🙏"
44)
45
46def filter_response(response: str, user_input: str) -> str:
47 # Remove URLs, markdown artifacts, and unwanted text
48 response = re.sub(r'http\S+', '', response)
49 response = re.sub(r'\[\w+\]|\(\w+\)|\*|\#', '', response)
50 response = response.split("http")[0].split("©")[0]
51
52 # Enforce brevity: Keep only the first two sentences
53 sentences = re.split(r'(?<=[.!?])\s+', response)
54 response = " ".join(sentences[:2]) # Keep only first 2 sentences
55
56 # Append crisis response if keywords detected
57 if any(keyword in user_input.lower() for keyword in CRISIS_KEYWORDS):
58 response += CRISIS_RESPONSE
59
60 # Correct grammar
61 response = correct_grammar(response)
62
63 return response
64
65def chat():
66
67 print("Chat with your fine-tuned Falcon model (type 'exit' to quit):")
68
69 system_instruction = (
70 "You are an empathetic AI specialized in mental health support. "
71 "Provide short, supportive, and comforting responses. "
72 "Validate the user's emotions and offer non-judgmental support. "
73 "If a crisis situation is detected, suggest reaching out to a mental health professional immediately. "
74 "Your responses should be clear, concise, and free from speculation. "
75 # "Examples:\n"
76 # "User: I feel really anxious lately.\n"
77 # "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"
78 # "User: I haven't been able to sleep well.\n"
79 # "AI: That sounds frustrating. Sleep troubles can be tough. Have you noticed anything that helps, like adjusting your bedtime routine?\n"
80 )
81
82 # Store short chat history for context
83 chat_history = []
84
85 while True:
86 user_input = input("\nYou: ")
87 if user_input.lower() == "exit":
88 break
89
90 # Maintain short chat history (last 2 exchanges)
91 chat_history.append(f"User: {user_input}")
92 chat_history = chat_history[-2:]
93
94 # Structure prompt
95 prompt = f"{system_instruction}\n" + "\n".join(chat_history) + "\nAI:"
96 inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
97
98 with torch.no_grad():
99 output = model.generate(
100 **inputs,
101 max_new_tokens=75,
102 pad_token_id=tokenizer.eos_token_id,
103 temperature=0.4,
104 top_p=0.9,
105 repetition_penalty=1.2,
106 do_sample=True,
107 no_repeat_ngram_size=2,
108 early_stopping=True
109 )
110
111 response = tokenizer.decode(output[0], skip_special_tokens=True).split("AI:")[-1].strip()
112 response = filter_response(response, user_input)
113 print(f"AI: {response}")
114
115 # Store AI response in history
116 chat_history.append(f"AI: {response}")
117
118chat()| Hyperparameter | Value |
|---|---|
| Precision | float16 |
| Optimizer | AdamW_32bit |
| Learning rate | 1.5e-4 |
| Weight decay | 1e-2 |
| Warmup Steps | 100 |
| Batch size | 2 |
| Training Epochs | 4 |
| Quantization | 8-Bit |