Create app.py

app.py

@@ -0,0 +1,178 @@
+import gradio as gr
+import torch
+import transformers
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from peft import PeftConfig, PeftModel
+import warnings
+from threading import Thread
+
+warnings.filterwarnings("ignore")
+
+PEFT_MODEL = "givyboy/phi-2-finetuned-mental-health-conversational"
+
+SYSTEM_PROMPT = """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.'."""
+
+USER_PROMPT = lambda x: f"""<HUMAN>: {x}\n<ASSISTANT>: """
+ADD_RESPONSE = lambda x, y: f"""<HUMAN>: {x}\n<ASSISTANT>: {y}"""
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.float16,
+)
+
+config = PeftConfig.from_pretrained(PEFT_MODEL)
+
+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,
+    offload_folder="offload/",
+    offload_state_dict=True,
+)
+
+peft_model = PeftModel.from_pretrained(
+    peft_base_model,
+    PEFT_MODEL,
+    offload_folder="offload/",
+    offload_state_dict=True,
+)
+peft_model = peft_model.to(DEVICE)
+
+peft_tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+peft_tokenizer.pad_token = peft_tokenizer.eos_token
+
+pipeline = transformers.pipeline(
+    "text-generation",
+    model=peft_model,
+    tokenizer=peft_tokenizer,
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
+    device_map="auto",
+)
+
+
+# def format_message(message: str, history: list[str], memory_limit: int = 3) -> str:
+#     if len(history) > memory_limit:
+#         history = history[-memory_limit:]
+
+#     if len(history) == 0:
+#         return f"{SYSTEM_PROMPT}\n{USER_PROMPT(message)}"
+
+#     formatted_message = f"{SYSTEM_PROMPT}\n{ADD_RESPONSE(history[0][0], history[0][1])}"
+
+#     for msg, ans in history[1:]:
+#         formatted_message += f"\n{ADD_RESPONSE(msg, ans)}"
+
+#     formatted_message += f"\n{USER_PROMPT(message)}"
+#     return formatted_message
+
+
+# def get_model_response(message: str, history: list[str]) -> str:
+#     formatted_message = format_message(message, history)
+#     sequences = pipeline(
+#         formatted_message,
+#         do_sample=True,
+#         top_k=10,
+#         num_return_sequences=1,
+#         eos_token_id=peft_tokenizer.eos_token_id,
+#         max_length=600,
+#     )[0]
+#     print(sequences["generated_text"])
+#     output = sequences["generated_text"].split("<ASSISTANT>:")[-1].strip()
+#     # print(f"Response: {output}")
+#     return output
+
+
+start_message = ""
+
+
+def user(message, history):
+    # Append the user's message to the conversation history
+    return "", history + [[message, ""]]
+
+
+def chat(message, history):
+    chat_history = []
+    for item in history:
+        chat_history.append({"role": "user", "content": item[0]})
+        if item[1] is not None:
+            chat_history.append({"role": "assistant", "content": item[1]})
+
+    message = f"{SYSTEM_PROMPT}\n{USER_PROMPT(message)}"
+    chat_history.append({"role": "user", "content": message})
+    messages = peft_tokenizer.apply_chat_template(chat_history, tokenize=False, add_generation_prompt=True)
+
+    # Tokenize the messages string
+    model_inputs = peft_tokenizer([messages], return_tensors="pt").to(DEVICE)
+    streamer = transformers.TextIteratorStreamer(
+        peft_tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
+    )
+    generate_kwargs = dict(
+        model_inputs,
+        streamer=streamer,
+        max_new_tokens=1024,
+        do_sample=True,
+        top_p=0.95,
+        top_k=1000,
+        temperature=0.75,
+        num_beams=1,
+    )
+    t = Thread(target=peft_model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    # Initialize an empty string to store the generated text
+    partial_text = ""
+    for new_text in streamer:
+        # print(new_text)
+        partial_text += new_text
+        # Yield an empty string to cleanup the message textbox and the updated conversation history
+        yield partial_text
+
+
+chat = gr.ChatInterface(fn=chat, title="Mental Health Chatbot - by Jayda Hunte")
+chat.launch(share=True)
+
+# import os
+# from openai import OpenAI
+# from dotenv import load_dotenv
+# import gradio as gr
+
+# load_dotenv()
+# API_KEY = os.getenv("OPENAI_API_KEY")
+# openai = OpenAI(api_key=API_KEY)
+
+# create_msg = lambda x, y: {"role": x, "content": y}
+
+# SYSTEM_PROMPT = create_msg(
+#     "system",
+#     """You are a helpful mental health chatbot, please answer with care. 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.'.""".strip(),
+# )
+
+
+# def predict(message, history):
+#     history_openai_format = []
+#     history_openai_format.append(SYSTEM_PROMPT)
+#     for human, assistant in history:
+#         history_openai_format.append({"role": "user", "content": human})
+#         history_openai_format.append({"role": "assistant", "content": assistant})
+#     history_openai_format.append({"role": "user", "content": message})
+
+#     response = openai.chat.completions.create(
+#         model="ft:gpt-3.5-turbo-0613:personal::8kBTG8eh", messages=history_openai_format, temperature=0.35, stream=True
+#     )
+
+#     partial_message = ""
+#     for chunk in response:
+#         if chunk.choices[0].delta.content is not None:
+#             partial_message = partial_message + chunk.choices[0].delta.content
+#             yield partial_message
+
+
+# gr.ChatInterface(fn=predict, title="Mental Health Chatbot").launch(share=True)