""" Star Coder2 chat demo """

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


# HF InferenceClient
client = InferenceClient("HuggingFaceH4/starchat2-15b-v0.1")


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
star_coder = gr.ChatInterface(
    chat,
    title="Star Coder 2",
    theme="soft",
    type="messages",
    description="Star Coder 2 is a advanced code assistant for software developers, "
    "built on top of the BigCode StarCoder 2 15B model.",
    additional_inputs=[
        gr.Textbox(
            value="You are a friendly coding assistant with the knowledge"
            "and experience of a senior software engineer.",
            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__":
    star_coder.launch()