Update app.py

app.py

@@ -1,125 +1,221 @@
-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
+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")
 
-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,
-  )
+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:
+#     <HUMAN>: question here
+#     <ASSISTANT>: answer here
+
+#     Begin!
+
+#     <HUMAN>: {query}
+#     <ASSISTANT>:"""
+
+#     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"<ASSISTANT>:.*?Begin!", bot_message, re.DOTALL)[1]
+#   except IndexError:
+#     pass
+
+#   bot_message = re.split(r'<ASSISTANT>\:?\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)
 
-  peft_model = PeftModel.from_pretrained(peft_base_model, PEFT_MODEL)
+# with gr.Blocks() as interface:
+#     gr.HTML("""<h1>Welcome to Mental Health Conversational AI</h1>""")
+#     gr.Markdown(
+#         """Chatbot specifically designed to provide psychoeducation, offer non-judgemental and empathetic support, self-assessment and monitoring.<br>
+#         Get instant response for any mental health related queries. If the chatbot seems you need external support, then it will respond appropriately.<br>"""
+#     )
 
-  peft_tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
-  peft_tokenizer.pad_token = peft_tokenizer.eos_token
+#     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")
 
-  return peft_model, peft_tokenizer
+#     llm_chain = init_llm_chain(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
+#     query.submit(user, [query, chatbot], [query, chatbot], queue=False).then(bot, chatbot, chatbot)
+#     clear.click(lambda: None, None, chatbot, queue=False)
 
-        @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:
-    <HUMAN>: question here
-    <ASSISTANT>: answer here
-
-    Begin!
-
-    <HUMAN>: {query}
-    <ASSISTANT>:"""
-
-    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"<ASSISTANT>:.*?Begin!", bot_message, re.DOTALL)[1]
-  except IndexError:
-    pass
-
-  bot_message = re.split(r'<ASSISTANT>\:?\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("""<h1>Welcome to Mental Health Conversational AI</h1>""")
-    gr.Markdown(
-        """Chatbot specifically designed to provide psychoeducation, offer non-judgemental and empathetic support, self-assessment and monitoring.<br>
-        Get instant response for any mental health related queries. If the chatbot seems you need external support, then it will respond appropriately.<br>"""
-    )
-
-    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()
+# interface.queue().launch()