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://fevhx93hdo4c770t.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 = 1512
    temperature = 0.9
    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|>", "|im_end|>"],
        #temperature=temperature,
        #top_p=top_p,
    ):
        token = message.choices[0].delta.content
        if not token:
            break
        response += token
        yield response.strip("|")


"""
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=["How can I lose weight safely?",
        "What is a healthy weight loss rate?",
        "How do I meal prep efficiently?",],
    # 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)",
    #     ),
    # ],
    css=".svelte-7ddecg h1, .message p {color: #64A149 !important; } .message{background-color: #64A149 !important; }"
)



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