Update app.py

app.py

@@ -1,67 +1,65 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 
-# Load model and tokenizer
-model_id = "Yhemhemoh/Gemma-2-2b-it-wazobia-wellness-bot"  # Replace with your model path
-model = AutoModelForCausalLM.from_pretrained(model_id)
-tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-def generate_response(user_input, selected_language="English"):
-    # Predefined instruction
-    predefined_instruction = (
-        "You are a highly skilled and empathetic mental health therapist fluent in "
-        "English, Yoruba, Igbo, and Hausa. Your task is to listen carefully to the user's concern "
-        "and respond with kindness, empathy, and respect. Always address the user directly in the second person, "
-        "avoiding third-person references. Never suggest harm to the user or others."
-    )
-    
-    # Format the input
-    formatted_input = (
-        f"<start_of_turn>user {predefined_instruction}\n\n"
-        f"User's concern:\n{user_input}\n\n"
-        f"Respond empathetically, kindly, and directly to the above concern in {selected_language}.<end_of_turn>\n"
-        f"<start_of_turn>model"
-    )
-    
-    # Tokenize
-    inputs = tokenizer(formatted_input, return_tensors='pt', padding=True, truncation=True)
-    
-    # Generate response
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_length=200,
-            num_return_sequences=1,
-            do_sample=True,
-            top_k=10,
-            top_p=0.8,
-            temperature=0.2,
-            no_repeat_ngram_size=3,
-        )
-    
-    # Decode and extract response
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response = text.split("model")[-1].strip()
-    
-    return response
+model = AutoModelForCausalLM.from_pretrained("hemhemoh/Gemma-2-2b-it-wazobia-wellness-bot")
 
-# Create Gradio interface
-iface = gr.Interface(
-    fn=generate_response,
-    inputs=[
-        gr.Textbox(label="Enter your concern", lines=3),
-        gr.Radio(["English", "Yoruba", "Igbo", "Hausa"], label="Select Language", value="English")
-    ],
-    outputs=gr.Textbox(label="Therapist's Response", lines=5),
-    title="Multilingual Mental Health Therapist",
-    description="Share your concerns and receive empathetic responses in English, Yoruba, Igbo, or Hausa.",
-    examples=[
-        ["I'm feeling very anxious about my upcoming exam", "English"],
-        ["Mo wa depressed, Awon obi mi ni kin kuro ni le tori mi o lowo", "Yoruba"],
-    ],
-    theme=gr.themes.Soft()
+
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,):
+    # Combine the system message with the user's input
+    message_prompt = "You are a highly skilled and empathetic mental health therapist fluent in English, Yoruba, Igbo, and Hausa. Respond to each user's concerns in the language they use to ensure comfort and understanding." + "\n\n"
+
+    # Add conversation history
+    prompt = ""
+    for user_input, assistant_response in history:
+        if user_input:
+             prompt = f"User's complaint:\n {user_input}\n"
+        if assistant_response:
+            prompt += f"assistant: {assistant_response}\n"
+
+    # Add the latest user message
+    prompt += f"{message_prompt}\nUser:{message}\n"
+
+    # Tokenize and generate response
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    #.to(device)
+    outputs = model.generate(
+        **inputs,
+        max_length=512,
+        temperature=0.2,
+        top_p=0.5,
+        no_repeat_ngram_size=3,)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # if "assistant:" in response:
+        # response = response.split("assistant:")[-1].strip()
+
+    yield response
+
+
+# Gradio ChatInterface with additional inputs
+demo = gr.ChatInterface(
+    respond,
+    # additional_inputs=[
+    #     # gr.Textbox(
+    #     #     value="Hi, How can we be of service to you, today?",
+    #     #     label="System message",
+    #     # ),
+    #     gr.Slider(minimum=1, maximum=512, value=512, step=1, label="Max new tokens"),
+    #     gr.Slider(minimum=0.1, maximum=1.0, value=0.1, step=0.1, label="Temperature"),
+    #     gr.Slider(
+    #         minimum=0.1,
+    #         maximum=1.0,
+    #         value=0.6,
+    #         step=0.05,
+    #         label="Top-p (nucleus sampling)",
+    #     ),
+    # ],
 )
 
-# Launch the app
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()