File size: 1,905 Bytes
9addef6
532bc0e
9addef6
9abe6ed
 
 
 
9addef6
 
9abe6ed
9addef6
 
532bc0e
9addef6
532bc0e
9addef6
 
 
 
 
 
 
9abe6ed
 
 
 
 
9addef6
9abe6ed
 
 
9addef6
9abe6ed
 
9addef6
9abe6ed
 
 
 
9addef6
9abe6ed
9addef6
 
 
 
 
 
 
532bc0e
9addef6
532bc0e
 
9addef6
 
 
 
 
 
 
9abe6ed
 
 
 
9addef6
9abe6ed
 
 
 
 
9addef6
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
""" 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}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})

        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    # 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 Coder2",
    theme="soft",
    description="Star Coder 2 is a code assistant trained on 15B parameters and 100K lines of code. "
    "It is a powerful code assistant that can help you write and debug code, as well as explain code to others.",
    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__":
    star_coder.launch()