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("unsloth/Llama-3.2-1B-Instruct")
client = InferenceClient( model="https://kjynd32snp9r6qb7.us-east-1.aws.endpoints.huggingface.cloud")

def respond(
    message,
    history: list[tuple[str, str]],
    # system_message,
    # max_tokens,
    # temperature,
    # top_p,
):
    
    system_message = "You are a Dietician Assistant specializing in providing general guidance on diet, "
    "nutrition, and healthy eating habits. Answer questions thoroughly with scientifically "
    "backed advice, practical tips, and easy-to-understand explanations. Keep in mind that "
    "your role is to assist, not replace a registered dietitian, so kindly remind users to "
    "consult a professional for personalized advice when necessary."
    max_tokens = 512
    temperature = 0.7
    top_p = 0.95
    
    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,
        stop=["<|im_end|><|im_end|>", "<|im_end|>"],
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        if not token:
            break
        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(
    fn=respond,
    #title="Hi there! I'm your Dietician Assistant, here to help you with general advice on diet, nutrition, and healthy eating habits. Let's explore your questions.",
    title="Your Personal Dietician Assistant: Expert Guidance on Healthy Eating and Nutrition",
    examples=["What are some healthy snacks I can eat if I feel hungry between meals?",
    "How can I improve my gut health through diet?",
    "Can you recommend some high-protein vegetarian meals?"]
    # additional_inputs=[
    #     gr.Textbox(value="You are a friendly Chatbot.", 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)",
    #     ),
    # ],
)



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