prathik31's picture
Update app.py
c3674d3
raw
history blame
4.86 kB
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")
# initialize and load PEFT model and tokenizer
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
# custom LLM chain to generate answer from PEFT model for each query
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 demo:
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)
demo.queue().launch(inline=False)