import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM

model_name = "llmaaz/mental_BART_model"  # Your model path

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

bart_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)

def respond(
    message,
    history: list[tuple[str, str]],
    max_tokens,
    temperature,
    top_p,
):
    """Generate a response from the fine-tuned BART model."""
    # Combine conversation history into a single input
    combined_input = ""
    for user_message, assistant_response in history:
        if user_message:
            combined_input += f"User: {user_message}\n"
        if assistant_response:
            combined_input += f"Assistant: {assistant_response}\n"

    # Add the new user message
    combined_input += f"User: {message}\nAssistant:"

    # Generate response using the BART pipeline
    response = bart_pipeline(
        combined_input, 
        max_length=max_tokens, 
        num_return_sequences=1, 
        temperature=temperature, 
        top_p=top_p
    )

    generated_text = response[0]["generated_text"].strip()
    yield generated_text

# Define the Gradio interface without system message
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        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)",
        ),
    ],
    title="Mental Health Assistant",
    description="This assistant uses a fine-tuned BART model to provide support for mental health discussions. Note: This is not a substitute for professional advice.",
)

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