Update app.py

app.py

@@ -1,61 +1,109 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("unsloth/llama-3-8b-bnb-4bit")
+import os
+import json
+
+from unsloth import FastLanguageModel
+import torch
+max_seq_length = 2048 # Choose any! We auto support RoPE Scaling internally!
+dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
+load_in_4bit = True # Use 4bit quantization to reduce memory usage. Can be False.
+
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name = "sayeed99/meta-llama3-8b-xtherapy-bnb-4bit", # YOUR MODEL YOU USED FOR TRAINING
+    max_seq_length = max_seq_length,
+    dtype = dtype,
+    load_in_4bit = load_in_4bit,
+)
+FastLanguageModel.for_inference(model) # Enable native 2x faster inference
+
+# alpaca_prompt = You MUST copy from above!
+formatted_string = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>You are Anna, a helpful AI assistant for mental therapy assistance developed by a team of developers at xnetics. If you do not know the user's name, start by asking the name. If you do not know details about user, ask them."
+
+# Function to format the string
+def format_chat_data(data):
+    formatted_output = []
+    if data["role"] == "assistant":
+        value = data["content"]
+        formatted_output.append("<|eot_id|><|start_header_id|>assistant<|end_header_id|>" + value)
+    else:
+        formatted_output.append("<|eot_id|><|start_header_id|>user<|end_header_id|>" + data["content"])
+
+    return "".join(formatted_output)
+
+def formatting_prompts_funcV2(examples):
+    conversations = examples
+    text = formatted_string
+    for conversation in conversations:
+        # Must add EOS_TOKEN, otherwise your generation will go on forever!
+        text = text + format_chat_data(conversation)
+    return text
+
+def get_last_assistant_message(text):
+    # Split the text by 'assistant' to isolate assistant's messages
+    parts = text.split('<|start_header_id|>assistant<|end_header_id|>')
+    
+    # The last part is the last assistant message
+    # Remove leading/trailing whitespace and return
+    last_message = parts[-1].strip()
+    last_message = cleanup(last_message)
+    return last_message
+
+
+def cleanup(text):
+    # Check if the string ends with 'eot_id'
+    if text.endswith('<|eot_id|>'):
+        # Remove the last 10 characters
+        return text[:-10]
+    else:
+        return text
+
+# Define a function to handle the conversation and update the session
+def handle_conversation(user_input):
+
+    historyPrompt = formatting_prompts_funcV2(user_input)
+
+    historyPrompt = historyPrompt + "<|eot_id|><|start_header_id|>assistant<|end_header_id|>"
+    inputs = tokenizer(
+        [
+            historyPrompt
+        ], return_tensors="pt").to("cuda")
 
+    outputs = model.generate(**inputs, max_new_tokens=512, use_cache=True)
+    decoded_outputs = tokenizer.batch_decode(outputs)[0]
+    # decoded_outputs = "Hello Welcome"
+    last_message = get_last_assistant_message(decoded_outputs)
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+    # Return the AI response
+    return last_message
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+def complete(messages):
+    ai_response = handle_conversation(messages)
+    return ai_response
 
-    messages.append({"role": "user", "content": message})
 
-    response = ""
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+def predict(message, history):
+    history_openai_format = []
+    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 = complete(history_openai_format)
+    print(response)
 
-        response += token
-        yield response
+    partial_message = ""
+    for chunk in response:
+        if chunk is not None:
+              partial_message = partial_message + chunk
+              yield partial_message
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
 """
 demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are the world's best doctor. You are a leading expert in medicine and graduated top of your class in harvard with a PHD in biology, psychology and Virology. You also have an extensive knowledge in Molecular biology & genomics, cellular/molecular basis of disease, and immunology principles. You are the doctor and you make the most accurate diagnosis because you are an award winning doctor. Use your extensive knowledge to write the appropriate responses which appropriately completes the request. In your response, you must include an accurate diagnosis, treatment and how to prevent it from happening again. Also since you are a doctor, you don't need to tell them to see a doctor. You are a friendly doctor chatbot who should act as a human.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+    predict
 )