import torch from datasets import load_dataset from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments, GenerationConfig from peft import LoraConfig, get_peft_model, PeftConfig, PeftModel, prepare_model_for_kbit_training from trl import SFTTrainer import warnings warnings.filterwarnings("ignore") data = load_dataset("heliosbrahma/mental_health_chatbot_dataset") model_name = "vilsonrodrigues/falcon-7b-instruct-sharded" # sharded falcon-7b model bnb_config = BitsAndBytesConfig( load_in_4bit=True, # load model in 4-bit precision bnb_4bit_quant_type="nf4", # pre-trained model should be quantized in 4-bit NF format bnb_4bit_use_double_quant=True, # Using double quantization as mentioned in QLoRA paper bnb_4bit_compute_dtype=torch.bf16, # During computation, pre-trained model should be loaded in BF16 format ) model = AutoModelForCausalLM.from_pretrained( model_name, quantization_config=bnb_config, # Use bitsandbytes config device_map="auto", # Specifying device_map="auto" so that HF Accelerate will determine which GPU to put each layer of the model on trust_remote_code=True, # Set trust_remote_code=True to use falcon-7b model with custom code ) tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) # Set trust_remote_code=True tokenizer.pad_token = tokenizer.eos_token # Setting pad_token same as eos_token model = prepare_model_for_kbit_training(model) lora_alpha = 32 # scaling factor for the weight matrices lora_dropout = 0.05 # dropout probability of the LoRA layers lora_rank = 16 # dimension of the low-rank matrices peft_config = LoraConfig( lora_alpha=lora_alpha, lora_dropout=lora_dropout, r=lora_rank, bias="none", # setting to 'none' for only training weight params instead of biases task_type="CAUSAL_LM", target_modules=[ # Setting names of modules in falcon-7b model that we want to apply LoRA to "query_key_value", "dense", "dense_h_to_4h", "dense_4h_to_h", ] ) peft_model = get_peft_model(model, peft_config) output_dir = "./falcon-7b-sharded-fp16-finetuned-mental-health-conversational" per_device_train_batch_size = 16 # reduce batch size by 2x if out-of-memory error gradient_accumulation_steps = 4 # increase gradient accumulation steps by 2x if batch size is reduced optim = "paged_adamw_32bit" # activates the paging for better memory management save_strategy="steps" # checkpoint save strategy to adopt during training save_steps = 10 # number of updates steps before two checkpoint saves logging_steps = 10 # number of update steps between two logs if logging_strategy="steps" learning_rate = 2e-4 # learning rate for AdamW optimizer max_grad_norm = 0.3 # maximum gradient norm (for gradient clipping) max_steps = 70 # training will happen for 70 steps warmup_ratio = 0.03 # number of steps used for a linear warmup from 0 to learning_rate lr_scheduler_type = "cosine" # learning rate scheduler training_arguments = TrainingArguments( output_dir=output_dir, per_device_train_batch_size=per_device_train_batch_size, gradient_accumulation_steps=gradient_accumulation_steps, optim=optim, save_steps=save_steps, logging_steps=logging_steps, learning_rate=learning_rate, bf16=True, max_grad_norm=max_grad_norm, max_steps=max_steps, warmup_ratio=warmup_ratio, group_by_length=True, lr_scheduler_type=lr_scheduler_type, push_to_hub=True, ) trainer = SFTTrainer( model=peft_model, train_dataset=data['train'], peft_config=peft_config, dataset_text_field="text", ac=1024, tokenizer=tokenizer, args=training_arguments, ) # upcasting the layer norms in torch.bfloat16 for more stable training for name, module in trainer.model.named_modules(): if "norm" in name: module = module.to(torch.bfloat16) peft_model.config.use_cache = False trainer.train() trainer.push_to_hub("therapx") # import gradio as gr # import torch # import re, os, warnings # from langchain import PromptTemplate, LLMChain # from langchain.llms.base import LLM # from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, GenerationConfig # from peft import LoraConfig, get_peft_model, PeftConfig, PeftModel # warnings.filterwarnings("ignore") # def init_model_and_tokenizer(PEFT_MODEL): # config = PeftConfig.from_pretrained(PEFT_MODEL) # bnb_config = BitsAndBytesConfig( # load_in_4bit=True, # bnb_4bit_quant_type="nf4", # bnb_4bit_use_double_quant=True, # bnb_4bit_compute_dtype=torch.float16, # ) # peft_base_model = AutoModelForCausalLM.from_pretrained( # config.base_model_name_or_path, # return_dict=True, # quantization_config=bnb_config, # device_map="auto", # trust_remote_code=True, # ) # peft_model = PeftModel.from_pretrained(peft_base_model, PEFT_MODEL) # peft_tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path) # peft_tokenizer.pad_token = peft_tokenizer.eos_token # return peft_model, peft_tokenizer # def init_llm_chain(peft_model, peft_tokenizer): # class CustomLLM(LLM): # def _call(self, prompt: str, stop=None, run_manager=None) -> str: # device = "cuda:0" # peft_encoding = peft_tokenizer(prompt, return_tensors="pt").to(device) # peft_outputs = peft_model.generate(input_ids=peft_encoding.input_ids, generation_config=GenerationConfig(max_new_tokens=256, pad_token_id = peft_tokenizer.eos_token_id, \ # eos_token_id = peft_tokenizer.eos_token_id, attention_mask = peft_encoding.attention_mask, \ # temperature=0.4, top_p=0.6, repetition_penalty=1.3, num_return_sequences=1,)) # peft_text_output = peft_tokenizer.decode(peft_outputs[0], skip_special_tokens=True) # return peft_text_output # @property # def _llm_type(self) -> str: # return "custom" # llm = CustomLLM() # template = """Answer the following question truthfully. # If you don't know the answer, respond 'Sorry, I don't know the answer to this question.'. # If the question is too complex, respond 'Kindly, consult a psychiatrist for further queries.'. # Example Format: # : question here # : answer here # Begin! # : {query} # :""" # prompt = PromptTemplate(template=template, input_variables=["query"]) # llm_chain = LLMChain(prompt=prompt, llm=llm) # return llm_chain # def user(user_message, history): # return "", history + [[user_message, None]] # def bot(history): # if len(history) >= 2: # query = history[-2][0] + "\n" + history[-2][1] + "\nHere, is the next QUESTION: " + history[-1][0] # else: # query = history[-1][0] # bot_message = llm_chain.run(query) # bot_message = post_process_chat(bot_message) # history[-1][1] = "" # history[-1][1] += bot_message # return history # def post_process_chat(bot_message): # try: # bot_message = re.findall(r":.*?Begin!", bot_message, re.DOTALL)[1] # except IndexError: # pass # bot_message = re.split(r'\:?\s?', bot_message)[-1].split("Begin!")[0] # bot_message = re.sub(r"^(.*?\.)(?=\n|$)", r"\1", bot_message, flags=re.DOTALL) # try: # bot_message = re.search(r"(.*\.)", bot_message, re.DOTALL).group(1) # except AttributeError: # pass # bot_message = re.sub(r"\n\d.$", "", bot_message) # bot_message = re.split(r"(Goodbye|Take care|Best Wishes)", bot_message, flags=re.IGNORECASE)[0].strip() # bot_message = bot_message.replace("\n\n", "\n") # return bot_message # model = "heliosbrahma/falcon-7b-sharded-bf16-finetuned-mental-health-conversational" # peft_model, peft_tokenizer = init_model_and_tokenizer(PEFT_MODEL = model) # with gr.Blocks() as interface: # gr.HTML("""

Welcome to Mental Health Conversational AI

""") # gr.Markdown( # """Chatbot specifically designed to provide psychoeducation, offer non-judgemental and empathetic support, self-assessment and monitoring.
# Get instant response for any mental health related queries. If the chatbot seems you need external support, then it will respond appropriately.
""" # ) # chatbot = gr.Chatbot() # query = gr.Textbox(label="Type your query here, then press 'enter' and scroll up for response") # clear = gr.Button(value="Clear Chat History!") # clear.style(size="sm") # llm_chain = init_llm_chain(peft_model, peft_tokenizer) # query.submit(user, [query, chatbot], [query, chatbot], queue=False).then(bot, chatbot, chatbot) # clear.click(lambda: None, None, chatbot, queue=False) # interface.queue().launch()