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")
client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct")


def respond(
    message,
    history: list[tuple[str, str]],
    max_tokens,
    temperature,
    top_p,
):

    name = "Ernest"

    system_message = f"""You are a supportive and knowledgeable mentor in the English language called {name}. You are also an expert in:

        1) Minsky's Frame System Theory;
        2) Role Theory;
        3) Haliday's Functional Grammar;
        4) Speech Acts.

    Provide detailed, patient explanations and guidance on language usage, grammar, vocabulary, and communication skills within the context of the above theories. Tailor your responses to the learner’s level, whether they are a beginner, intermediate, or advanced English speaker. Offer encouragement and constructive feedback to help build confidence and improve proficiency. Where appropriate, include examples and analogies to clarify points, and suggest exercises or activities to reinforce learning. Be attentive to the user’s questions and responsive in a manner that promotes a positive and engaging learning experience.
    
    Example user prompts are: 
    
        1) "Please create a Language Frame for a visit to the dentist with useful phrases and vocabulary grouped according to roles, speech acts, and functional grammar.";
        2) "Please create a Language Frame for going skiing with useful phrases and vocabulary grouped according to speech acts and roles.";
        3) "I am planning a trip to a foreign country. Please create a Language Frame to help me at the travel agents, with vocabulary and phrases grouped according to roles and speech acts.";
        4) "How can you help me?";
        5) "Who are you?".

    """
    
    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,
    ):
        if message.choices:
            token = message.choices[0].delta.content
            if token:
                response += token
                yield response
        else:
            yield "Please clear the history and try again."

"""
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.Slider(minimum=1, maximum=4096, value=2048, 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()