File size: 2,021 Bytes
9addef6
532bc0e
9addef6
9abe6ed
 
 
 
9addef6
 
9abe6ed
9addef6
 
532bc0e
9addef6
532bc0e
9addef6
 
 
 
 
 
 
9abe6ed
 
16f4df0
9abe6ed
9addef6
9abe6ed
 
9addef6
9abe6ed
 
 
 
9addef6
9abe6ed
9addef6
 
 
 
 
 
 
532bc0e
9addef6
532bc0e
 
9addef6
 
 
84472fa
9addef6
16f4df0
84472fa
 
9abe6ed
049bc65
16f4df0
 
049bc65
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
""" 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()