jakewatson commited on
Commit
04efeb6
·
1 Parent(s): 2cbcf79
Files changed (1) hide show
  1. app.py +135 -43
app.py CHANGED
@@ -1,63 +1,155 @@
1
  import gradio as gr
2
  from huggingface_hub import InferenceClient
 
 
3
 
4
- """
5
- For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
6
- """
7
  client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
8
 
 
 
9
 
10
  def respond(
11
  message,
12
  history: list[tuple[str, str]],
13
- system_message,
14
- max_tokens,
15
- temperature,
16
- top_p,
 
17
  ):
18
- messages = [{"role": "system", "content": system_message}]
 
19
 
20
- for val in history:
21
- if val[0]:
22
- messages.append({"role": "user", "content": val[0]})
23
- if val[1]:
24
- messages.append({"role": "assistant", "content": val[1]})
25
 
26
- messages.append({"role": "user", "content": message})
 
 
 
 
 
 
 
 
27
 
28
- response = ""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
29
 
30
- for message in client.chat_completion(
31
- messages,
32
- max_tokens=max_tokens,
33
- stream=True,
34
- temperature=temperature,
35
- top_p=top_p,
36
- ):
37
- token = message.choices[0].delta.content
 
38
 
39
- response += token
40
- yield response
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
41
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
42
  """
43
- For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
44
- """
45
- demo = gr.ChatInterface(
46
- respond,
47
- additional_inputs=[
48
- gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
49
- gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
50
- gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
51
- gr.Slider(
52
- minimum=0.1,
53
- maximum=1.0,
54
- value=0.95,
55
- step=0.05,
56
- label="Top-p (nucleus sampling)",
57
- ),
58
- ],
59
- )
60
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
61
 
62
  if __name__ == "__main__":
63
- demo.launch()
 
1
  import gradio as gr
2
  from huggingface_hub import InferenceClient
3
+ import torch
4
+ from transformers import pipeline
5
 
6
+ # Inference client setup
 
 
7
  client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
8
+ pipe = pipeline("text-generation", "microsoft/Phi-3-mini-4k-instruct", torch_dtype=torch.bfloat16, device_map="auto")
9
 
10
+ # Global flag to handle cancellation
11
+ stop_inference = False
12
 
13
  def respond(
14
  message,
15
  history: list[tuple[str, str]],
16
+ system_message="You are a friendly Chatbot.",
17
+ max_tokens=512,
18
+ temperature=0.7,
19
+ top_p=0.95,
20
+ use_local_model=False,
21
  ):
22
+ global stop_inference
23
+ stop_inference = False # Reset cancellation flag
24
 
25
+ # Initialize history if it's None
26
+ if history is None:
27
+ history = []
 
 
28
 
29
+ if use_local_model:
30
+ # local inference
31
+ messages = [{"role": "system", "content": system_message}]
32
+ for val in history:
33
+ if val[0]:
34
+ messages.append({"role": "user", "content": val[0]})
35
+ if val[1]:
36
+ messages.append({"role": "assistant", "content": val[1]})
37
+ messages.append({"role": "user", "content": message})
38
 
39
+ response = ""
40
+ for output in pipe(
41
+ messages,
42
+ max_new_tokens=max_tokens,
43
+ temperature=temperature,
44
+ do_sample=True,
45
+ top_p=top_p,
46
+ ):
47
+ if stop_inference:
48
+ response = "Inference cancelled."
49
+ yield history + [(message, response)]
50
+ return
51
+ token = output['generated_text'][-1]['content']
52
+ response += token
53
+ yield history + [(message, response)] # Yield history + new response
54
 
55
+ else:
56
+ # API-based inference
57
+ messages = [{"role": "system", "content": system_message}]
58
+ for val in history:
59
+ if val[0]:
60
+ messages.append({"role": "user", "content": val[0]})
61
+ if val[1]:
62
+ messages.append({"role": "assistant", "content": val[1]})
63
+ messages.append({"role": "user", "content": message})
64
 
65
+ response = ""
66
+ for message_chunk in client.chat_completion(
67
+ messages,
68
+ max_tokens=max_tokens,
69
+ stream=True,
70
+ temperature=temperature,
71
+ top_p=top_p,
72
+ ):
73
+ if stop_inference:
74
+ response = "Inference cancelled."
75
+ yield history + [(message, response)]
76
+ return
77
+ if stop_inference:
78
+ response = "Inference cancelled."
79
+ break
80
+ token = message_chunk.choices[0].delta.content
81
+ response += token
82
+ yield history + [(message, response)] # Yield history + new response
83
 
84
+
85
+ def cancel_inference():
86
+ global stop_inference
87
+ stop_inference = True
88
+
89
+ # Custom CSS for a fancy look
90
+ custom_css = """
91
+ #main-container {
92
+ background-color: #f0f0f0;
93
+ font-family: 'Arial', sans-serif;
94
+ }
95
+ .gradio-container {
96
+ max-width: 700px;
97
+ margin: 0 auto;
98
+ padding: 20px;
99
+ background: white;
100
+ box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
101
+ border-radius: 10px;
102
+ }
103
+ .gr-button {
104
+ background-color: #4CAF50;
105
+ color: white;
106
+ border: none;
107
+ border-radius: 5px;
108
+ padding: 10px 20px;
109
+ cursor: pointer;
110
+ transition: background-color 0.3s ease;
111
+ }
112
+ .gr-button:hover {
113
+ background-color: #45a049;
114
+ }
115
+ .gr-slider input {
116
+ color: #4CAF50;
117
+ }
118
+ .gr-chat {
119
+ font-size: 16px;
120
+ }
121
+ #title {
122
+ text-align: center;
123
+ font-size: 2em;
124
+ margin-bottom: 20px;
125
+ color: #333;
126
+ }
127
  """
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
128
 
129
+ # Define the interface
130
+ with gr.Blocks(css=custom_css) as demo:
131
+ gr.Markdown("<h1 style='text-align: center;'>🌟 Fancy AI Chatbot 🌟</h1>")
132
+ gr.Markdown("Interact with the AI chatbot using customizable settings below.")
133
+
134
+ with gr.Row():
135
+ system_message = gr.Textbox(value="You are a friendly Chatbot.", label="System message", interactive=True)
136
+ use_local_model = gr.Checkbox(label="Use Local Model", value=False)
137
+
138
+ with gr.Row():
139
+ max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
140
+ temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
141
+ top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
142
+
143
+ chat_history = gr.Chatbot(label="Chat")
144
+
145
+ user_input = gr.Textbox(show_label=False, placeholder="Type your message here...")
146
+
147
+ cancel_button = gr.Button("Cancel Inference", variant="danger")
148
+
149
+ # Adjusted to ensure history is maintained and passed correctly
150
+ user_input.submit(respond, [user_input, chat_history, system_message, max_tokens, temperature, top_p, use_local_model], chat_history)
151
+
152
+ cancel_button.click(cancel_inference)
153
 
154
  if __name__ == "__main__":
155
+ demo.launch(share=False) # Remove share=True because it's not supported on HF Spaces