File size: 4,886 Bytes
1694b02
2831539
 
 
1694b02
2831539
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1694b02
99f28a3
2831539
 
 
99f28a3
2831539
1694b02
2831539
 
 
 
 
 
 
 
 
 
 
 
 
 
1694b02
 
2831539
 
 
 
1694b02
2831539
 
 
 
 
 
99f28a3
2831539
 
 
99f28a3
 
2831539
99f28a3
2831539
 
 
99f28a3
 
 
2831539
 
 
99f28a3
 
 
 
 
2831539
 
 
99f28a3
2831539
99f28a3
2831539
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1694b02
2831539
1694b02
 
2831539
 
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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
import gradio as gr
import uuid
import json
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the DeepSeek-R1 model
model = AutoModelForCausalLM.from_pretrained(
    "deepseek-ai/DeepSeek-R1",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1")

def format_messages_for_model(messages):
    """
    Format the conversation history into a single string compatible with the DeepSeek-R1 model.
    """
    formatted_messages = ""
    for msg in messages:
        if msg["role"] == "user":
            formatted_messages += f"User: {msg['content']}\n"
        elif msg["role"] == "assistant":
            formatted_messages += f"DeepSeek-R1: {msg['content']}\n"
    return formatted_messages

def generate_response(user_input, messages):
    """
    Generate a response using DeepSeek-R1 and return it as a string.
    """
    formatted_history = format_messages_for_model(messages)
    full_input = f"{formatted_history}User: {user_input}\nDeepSeek-R1: "

    inputs = tokenizer(full_input, return_tensors="np")
    outputs = model.generate(
        **inputs,
        max_length=2048,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        return_dict_in_generate=True,
    )
    
    response = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
    return response.split("DeepSeek-R1: ")[-1].strip()

def chatbot(
    message,
    history: list[tuple[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
):
    """
    Main function for generating responses and saving chat history.
    """
    # Generate a unique session ID for this chat session
    session_id = str(uuid.uuid4())

    # Initialize the messages with the system message
    messages = []
    if system_message:
        messages.append({"role": "system", "content": system_message})

    # Add previous chat history
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})

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

    # Generate response
    response = generate_response(message, messages)

    # Add the bot's response to messages
    messages.append({"role": "assistant", "content": response})

    # Update chat history in the interface
    new_history = history + [(message, response)]

    # Save chat history to Hugging Face storage
    storage = gr.HuggingFace()
    storage.set(key=session_id, value={"messages": messages})

    return response, new_history

def create_chat_interface():
    """
    Create the Gradio chat interface.
    """
    with gr.Blocks() as app:
        session_id = str(uuid.uuid4())  # Generate session ID at startup

        with gr.Row():
            chatbot_interface = gr.Chatbot(label="DeepSeek-R1 Chatbot")
            input_box = gr.Textbox(
                label="Type your message here",
                placeholder="Type your message here...",
                show_label=True,
            )

        # Button to submit the message
        submit_button = gr.Button("Send")

        # Store chat history in the interface
        history = gr.State([])

        # Define the system message
        system_message = gr.Textbox(
            value="You are DeepSeek-R1, an AI assistant.",
            label="System Message",
            show_label=True,
            visible=True,
        )

        # Define inference parameters
        with gr.Accordion("Parameters", open=True):
            max_tokens = gr.Slider(
                minimum=1,
                maximum=2048,
                value=512,
                step=1,
                label="Max Tokens",
                info="Maximum length of the generated response.",
            )
            temperature = gr.Slider(
                minimum=0.0,
                maximum=2.0,
                value=1.0,
                step=0.1,
                label="Temperature",
                info="Higher values produce more diverse outputs.",
            )
            top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.9,
                step=0.05,
                label="Top-p",
                info="Controls diversity through nucleus sampling.",
            )

        # Define the submit function
        submit_button.click(
            fn=chatbot,
            inputs=[
                input_box,
                history,
                system_message,
                max_tokens,
                temperature,
                top_p,
            ],
            outputs=[gr.Textbox(), chatbot_interface],
            api_name="chatbot",
        )

    return app

if __name__ == "__main__":
    app = create_chat_interface()
    app.launch()