from huggingface_hub import InferenceClient
import gradio as gr

# Upgraded to Mistral-7B-v0.3 
client = InferenceClient(
    "mistralai/Mistral-7B-Instruct-v0.3"
)


def format_prompt(message, history):
  prompt = "<s>"
  for user_prompt, bot_response in history:
    prompt += f"[INST] {user_prompt} [/INST]"
    prompt += f" {bot_response}</s> "
  prompt += f"[INST] {message} [/INST]"
  return prompt

def generate(
    prompt, history, temperature=0.9, max_new_tokens=900, top_p=0.95, repetition_penalty=1.0,
):
    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True,
        seed=42,
    )

    formatted_prompt = format_prompt(prompt, history)

    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""

    for response in stream:
        output += response.token.text
        yield output
    return output


additional_inputs=[
    gr.Slider(
        label="Temperature",
        value=0.9,
        minimum=0.0,
        maximum=1.0,
        step=0.05,
        interactive=True,
        info="Higher values produce more diverse outputs",
    ),
    gr.Slider(
        label="Max new tokens",
        value=900,
        minimum=0,
        maximum=1048,
        step=64,
        interactive=True,
        info="The maximum numbers of new tokens",
    ),
    gr.Slider(
        label="Top-p (nucleus sampling)",
        value=0.90,
        minimum=0.0,
        maximum=1,
        step=0.05,
        interactive=True,
        info="Higher values sample more low-probability tokens",
    ),
    gr.Slider(
        label="Repetition penalty",
        value=1.2,
        minimum=1.0,
        maximum=2.0,
        step=0.05,
        interactive=True,
        info="Penalize repeated tokens",
    )
]

css = """
  #mkd {
    height: 500px; 
    overflow: auto; 
    border: 1px solid #ccc; 
  }
"""

with gr.Blocks(css=css) as ai_chat:
    gr.HTML("<h1><center>AI Conversation<h1><center>")
    gr.HTML("<h3><center>How can I help you? You can converse with me and say more💬<h3><center>")
    gr.HTML("<h3><center>To try, select a prompt from below and hit submit<h3><center>")
    gr.HTML("<h3><center>Have a wonderful day! 📚<h3><center>")
    gr.ChatInterface(
        generate,
        additional_inputs=additional_inputs,
        examples=[["List fun activities in Boston."], ["How to spend a weekend in San Francisco?"], ["What is the secret to life?"], ["Write me a recipe for a quick vegetarian breakfast."],["What is the future for full stack engineers?"], 
                 ["Create a complete plan for daily healthy habbits."], ["What is optogenetic simulation?"], ["How to conduct a neuroscience experiment using holography?"], ["What is non-invasive brain stimulation?"], ["Tell me lifestyle of people living in Auckland, NZ"], ["Make a tour plan for Los Angeles metro area."]]
    )
'''
By enabling the queue you can control when users know their position in the queue, and set a limit on maximum number of events allowed.
        Parameters:
            status_update_rate: If "auto", Queue will send status estimations to all clients whenever a job is finished. Otherwise Queue will send status at regular intervals set by this parameter as the number of seconds.
            api_open: If True, the REST routes of the backend will be open, allowing requests made directly to those endpoints to skip the queue.
            max_size: The maximum number of events the queue will store at any given moment. If the queue is full, new events will not be added and a user will receive a message saying that the queue is full. If None, the queue size will be unlimited.
            concurrency_count: Deprecated. Set the concurrency_limit directly on event listeners e.g. btn.click(fn, ..., concurrency_limit=10) or gr.Interface(concurrency_limit=10). If necessary, the total number of workers can be configured via `max_threads` in launch().
            default_concurrency_limit: The default value of `concurrency_limit` to use for event listeners that don't specify a value. Can be set by environment variable GRADIO_DEFAULT_CONCURRENCY_LIMIT. Defaults to 1 if not set otherwise.
    replace deprecated concurency_count to concurrency_limit
'''
#ai_chat.queue(concurrency_limit=None, max_size=250).launch(debug=True)
ai_chat.queue(max_size=250).launch(debug=True)