""" Star Coder2 chat demo """

from typing import List, Tuple, Union
import gradio as gr
from huggingface_hub import InferenceClient


# HF InferenceClient
client = InferenceClient("microsoft/Phi-3.5-mini-instruct")


def chat(
    message: str,
    history: List[Tuple[str, str]],
    system_message: str,
    max_tokens: Union[int, None],
    temperature: Union[float, None],
    top_p: Union[float, None],
):
    """Code assistant"""

    # Chat history
    messages = [{"role": "system", "content": system_message}]

    messages.extend(history)

    # Add user message
    messages.append({"role": "user", "content": message})

    llm_message = client.chat_completion(
        messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    )

    # Add chatbot message
    messages.append(
        {
            "role": "assistant",
            "content": llm_message.choices[0].message.content,
        }
    )

    yield llm_message.choices[0].message.content


# UI
demo = gr.ChatInterface(
    chat,
    theme="soft",
    type="messages",
    title="Phi-3.5-mini-instruct",
    description="Phi-3.5-mini is a lightweight, state-of-the-art open model built upon "
    "datasets used for Phi-3 - synthetic data and filtered publicly available websites - "
    "with a focus on very high-quality, reasoning dense data.",
    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",
        ),
    ],
)


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