import gradio as gr
import os
import pandas as pd
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import InferenceClient

# Load the medical dataset from GitHub
url = "https://raw.githubusercontent.com/wifix199/health_dataset/refs/heads/main/dataset.csv"
data = pd.read_csv(url)

# Assuming the dataset has 'Symptom' and 'Disease' columns
symptoms = data['symptoms'].tolist()
diseases = data['disease'].tolist()
HF_TOKEN = os.getenv("HF_TOKEN")

# Load an open-source LLM (e.g., facebook/opt-350m)
model_name = "mistralai/Mistral-7B-Instruct-v0.3"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Create a text generation pipeline with 'max_new_tokens' to limit output length
text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Define the chatbot logic
async def chatbot(messages, symptom_input):
    # Add user input to the chat history
    messages.append(("User", symptom_input))
    
    # Find the closest matching symptom in the dataset
    matched_diseases = [disease for symptom, disease in zip(symptoms, diseases) if symptom_input.lower() in symptom.lower()]
    
    if matched_diseases:
        response = f"Based on the symptom '{symptom_input}', the potential disease(s) might be: {', '.join(matched_diseases)}."
    else:
        # Generate a more refined response using the LLM if no match is found
        refined_prompt = (
            f"The user has reported a symptom {symptom_input}."
            "Please provide a brief and informative explanation, focusing on possible causes, and avoid repeating the user's input. "
            "Ensure the response is concise and helpful."
        )
        
        try:
            generated = text_generator(refined_prompt, max_new_tokens=150, num_return_sequences=1)
            response = generated[0]['generated_text'].strip()
        except Exception as e:
            response = f"Error during text generation: {str(e)}"
    
    # Add the bot's response to the chat history
    messages.append(("Bot", response))
    return messages, ""

# Define the TITLE
TITLE = "<h1>Healthcare Chatbot</h1>"

# Create the Gradio interface with advanced settings
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.HTML(TITLE)
    
    chatbot_ui = gr.Chatbot(label="Healthcare Assistant")
    
    # Input field for symptom input
    symptom_input = gr.Textbox(placeholder="Enter your symptom...", show_label=False)
    submit_button = gr.Button("Submit")

    # Parameters for chatbot fine-tuning
    temperature = gr.Slider(
        minimum=0,
        maximum=1,
        step=0.1,
        value=0.4,
        label="Temperature",
        render=False
    )
    max_tokens = gr.Slider(
        minimum=100,
        maximum=500,
        step=10,
        value=300,
        label="Max new tokens",
        render=False
    )
    top_p = gr.Slider(
        minimum=0.0,
        maximum=1.0,
        step=0.01,
        value=0.70,
        label="Top-p",
        render=False
    )
    top_k = gr.Slider(
        minimum=1,
        maximum=50,
        step=1,
        value=10,
        label="Top-k",
        render=False
    )
    repetition_penalty = gr.Slider(
        minimum=0.0,
        maximum=2.0,
        step=0.1,
        value=1.1,
        label="Repetition penalty",
        render=False
    )

    # Maintain chat history
    chat_history = gr.State([])

    # Define the interaction between input and chatbot
    submit_button.click(fn=chatbot, inputs=[chat_history, symptom_input], outputs=[chatbot_ui, symptom_input])

# Launch the chatbot on the specified IP address and port
demo.launch(server_name="10.172.2.37", server_port=5001)


Can you change the way how the bot respond