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("HuggingFaceH4/zephyr-7b-beta")


def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    system_message = "The use of AI, inclusive of Generative AI, in healthcare is evolving rapidly and has the 'potential to enhance healthcare outcomes by improving clinical trials, medical diagnosis and treatment, self-management of care, and personalized care. AI has been a hot topic and has captured a considerable amount of attention due to the recent advancements and implementation of this type of technology. When it comes to healthcare, AI is already actively being used in this industry on a smaller level, but there are various factors that prevent large-scale automation which we will further explore."
    messages = [{"role": "system", "content": system_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]})

    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

        response += token
        yield response

"""
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 = "The use of AI, inclusive of Generative AI, in healthcare is evolving rapidly and has the 'potential to enhance healthcare outcomes by improving clinical trials, medical diagnosis and treatment, self-management of care, and personalized care. AI has been a hot topic and has captured a considerable amount of attention due to the recent advancements and implementation of this type of technology. When it comes to healthcare, AI is already actively being used in this industry on a smaller level, but there are various factors that prevent large-scale automation which we will further explore.", 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)",
        ),
    ],

    examples = [ 
        ["What are the primary applications of AI in healthcare?"],
        ["How does AI improve diagnostic accuracy?"],
        ["How can AI help in early disease detection?"]
    ],
    title = 'AI in Healthcare'
)


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