import gradio as gr
from huggingface_hub import InferenceClient

# Instantiate the client with the mental health model or any suitable model.
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    # Define the system message to guide the assistant's behavior
    messages = [{"role": "system", "content": system_message}]

    # Add history of conversation to maintain context
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    # Add the user's current message to the conversation context
    messages.append({"role": "user", "content": message})

    response = ""

    # Stream the response token by token
    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
        response += token
        yield response

# Customize the system message for mental health
default_system_message = """
You are a compassionate mental health specialist trained to listen empathetically and offer support. 
When engaging with users, make sure to respond with kindness and provide general emotional support. 
Avoid giving specific medical or clinical advice, but offer guidance, validate feelings, and suggest appropriate resources when needed. 
Encourage open conversations and create a safe, non-judgmental space for the user to share.
"""

# Create the Gradio interface with additional parameters for user configuration
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value=default_system_message, label="System Message (Mental Health Specialist)"),
        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)",
        ),
    ],
)

if __name__ == "__main__":
    demo.launch()