forked from h2oai/h2ogpt
-
Notifications
You must be signed in to change notification settings - Fork 0
/
client_test.py
337 lines (264 loc) · 12.7 KB
/
client_test.py
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
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
"""
Client test.
Run server:
python generate.py --base_model=h2oai/h2ogpt-oig-oasst1-512-6_9b
NOTE: For private models, add --use-auth_token=True
NOTE: --infer_devices=True (default) must be used for multi-GPU in case see failures with cuda:x cuda:y mismatches.
Currently, this will force model to be on a single GPU.
Then run this client as:
python client_test.py
For HF spaces:
HOST="https://h2oai-h2ogpt-chatbot.hf.space" python client_test.py
Result:
Loaded as API: https://h2oai-h2ogpt-chatbot.hf.space ✔
{'instruction_nochat': 'Who are you?', 'iinput_nochat': '', 'response': 'I am h2oGPT, a large language model developed by LAION.', 'sources': ''}
For demo:
HOST="https://gpt.h2o.ai" python client_test.py
Result:
Loaded as API: https://gpt.h2o.ai ✔
{'instruction_nochat': 'Who are you?', 'iinput_nochat': '', 'response': 'I am h2oGPT, a chatbot created by LAION.', 'sources': ''}
NOTE: Raw output from API for nochat case is a string of a python dict and will remain so if other entries are added to dict:
{'response': "I'm h2oGPT, a large language model by H2O.ai, the visionary leader in democratizing AI.", 'sources': ''}
"""
import ast
import time
import os
import markdown # pip install markdown
import pytest
from bs4 import BeautifulSoup # pip install beautifulsoup4
from enums import DocumentChoices
debug = False
os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
def get_client(serialize=True):
from gradio_client import Client
client = Client(os.getenv('HOST', "http://localhost:7860"), serialize=serialize)
if debug:
print(client.view_api(all_endpoints=True))
return client
def get_args(prompt, prompt_type, chat=False, stream_output=False,
max_new_tokens=50,
top_k_docs=3,
langchain_mode='Disabled', prompt_dict=None):
from collections import OrderedDict
kwargs = OrderedDict(instruction=prompt if chat else '', # only for chat=True
iinput='', # only for chat=True
context='',
# streaming output is supported, loops over and outputs each generation in streaming mode
# but leave stream_output=False for simple input/output mode
stream_output=stream_output,
prompt_type=prompt_type,
prompt_dict=prompt_dict,
temperature=0.1,
top_p=0.75,
top_k=40,
num_beams=1,
max_new_tokens=max_new_tokens,
min_new_tokens=0,
early_stopping=False,
max_time=20,
repetition_penalty=1.0,
num_return_sequences=1,
do_sample=True,
chat=chat,
instruction_nochat=prompt if not chat else '',
iinput_nochat='', # only for chat=False
langchain_mode=langchain_mode,
top_k_docs=top_k_docs,
chunk=True,
chunk_size=512,
document_choice=[DocumentChoices.All_Relevant.name],
)
from generate import eval_func_param_names
assert len(set(eval_func_param_names).difference(set(list(kwargs.keys())))) == 0
if chat:
# add chatbot output on end. Assumes serialize=False
kwargs.update(dict(chatbot=[]))
return kwargs, list(kwargs.values())
@pytest.mark.skip(reason="For manual use against some server, no server launched")
def test_client_basic(prompt_type='human_bot'):
return run_client_nochat(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
def run_client_nochat(prompt, prompt_type, max_new_tokens):
kwargs, args = get_args(prompt, prompt_type, chat=False, max_new_tokens=max_new_tokens)
api_name = '/submit_nochat'
client = get_client(serialize=True)
res = client.predict(
*tuple(args),
api_name=api_name,
)
print("Raw client result: %s" % res, flush=True)
res_dict = dict(prompt=kwargs['instruction_nochat'], iinput=kwargs['iinput_nochat'],
response=md_to_text(res))
print(res_dict)
return res_dict, client
@pytest.mark.skip(reason="For manual use against some server, no server launched")
def test_client_basic_api(prompt_type='human_bot'):
return run_client_nochat_api(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
def run_client_nochat_api(prompt, prompt_type, max_new_tokens):
kwargs, args = get_args(prompt, prompt_type, chat=False, max_new_tokens=max_new_tokens)
api_name = '/submit_nochat_api' # NOTE: like submit_nochat but stable API for string dict passing
client = get_client(serialize=True)
res = client.predict(
str(dict(kwargs)),
api_name=api_name,
)
print("Raw client result: %s" % res, flush=True)
res_dict = dict(prompt=kwargs['instruction_nochat'], iinput=kwargs['iinput_nochat'],
response=md_to_text(ast.literal_eval(res)['response']),
sources=ast.literal_eval(res)['sources'])
print(res_dict)
return res_dict, client
@pytest.mark.skip(reason="For manual use against some server, no server launched")
def test_client_basic_api_lean(prompt_type='human_bot'):
return run_client_nochat_api_lean(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
def run_client_nochat_api_lean(prompt, prompt_type, max_new_tokens):
kwargs = dict(instruction_nochat=prompt)
api_name = '/submit_nochat_api' # NOTE: like submit_nochat but stable API for string dict passing
client = get_client(serialize=True)
res = client.predict(
str(dict(kwargs)),
api_name=api_name,
)
print("Raw client result: %s" % res, flush=True)
res_dict = dict(prompt=kwargs['instruction_nochat'],
response=md_to_text(ast.literal_eval(res)['response']),
sources=ast.literal_eval(res)['sources'])
print(res_dict)
return res_dict, client
@pytest.mark.skip(reason="For manual use against some server, no server launched")
def test_client_basic_api_lean_morestuff(prompt_type='human_bot'):
return run_client_nochat_api_lean_morestuff(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
def run_client_nochat_api_lean_morestuff(prompt, prompt_type='human_bot', max_new_tokens=512):
kwargs = dict(
instruction='',
iinput='',
context='',
stream_output=False,
prompt_type=prompt_type,
temperature=0.1,
top_p=0.75,
top_k=40,
num_beams=1,
max_new_tokens=256,
min_new_tokens=0,
early_stopping=False,
max_time=20,
repetition_penalty=1.0,
num_return_sequences=1,
do_sample=True,
chat=False,
instruction_nochat=prompt,
iinput_nochat='',
langchain_mode='Disabled',
top_k_docs=4,
document_choice=['All'],
)
api_name = '/submit_nochat_api' # NOTE: like submit_nochat but stable API for string dict passing
client = get_client(serialize=True)
res = client.predict(
str(dict(kwargs)),
api_name=api_name,
)
print("Raw client result: %s" % res, flush=True)
res_dict = dict(prompt=kwargs['instruction_nochat'],
response=md_to_text(ast.literal_eval(res)['response']),
sources=ast.literal_eval(res)['sources'])
print(res_dict)
return res_dict, client
@pytest.mark.skip(reason="For manual use against some server, no server launched")
def test_client_chat(prompt_type='human_bot'):
return run_client_chat(prompt='Who are you?', prompt_type=prompt_type, stream_output=False, max_new_tokens=50,
langchain_mode='Disabled')
@pytest.mark.skip(reason="For manual use against some server, no server launched")
def test_client_chat_stream(prompt_type='human_bot'):
return run_client_chat(prompt="Tell a very long kid's story about birds.", prompt_type=prompt_type,
stream_output=True, max_new_tokens=512,
langchain_mode='Disabled')
def run_client_chat(prompt, prompt_type, stream_output, max_new_tokens, langchain_mode, prompt_dict=None):
client = get_client(serialize=False)
kwargs, args = get_args(prompt, prompt_type, chat=True, stream_output=stream_output,
max_new_tokens=max_new_tokens, langchain_mode=langchain_mode,
prompt_dict=prompt_dict)
return run_client(client, prompt, args, kwargs)
def run_client(client, prompt, args, kwargs, do_md_to_text=True, verbose=False):
assert kwargs['chat'], "Chat mode only"
res = client.predict(*tuple(args), api_name='/instruction')
args[-1] += [res[-1]]
res_dict = kwargs
res_dict['prompt'] = prompt
if not kwargs['stream_output']:
res = client.predict(*tuple(args), api_name='/instruction_bot')
res_dict['response'] = res[0][-1][1]
print(md_to_text(res_dict['response'], do_md_to_text=do_md_to_text))
return res_dict, client
else:
job = client.submit(*tuple(args), api_name='/instruction_bot')
res1 = ''
while not job.done():
outputs_list = job.communicator.job.outputs
if outputs_list:
res = job.communicator.job.outputs[-1]
res1 = res[0][-1][-1]
res1 = md_to_text(res1, do_md_to_text=do_md_to_text)
print(res1)
time.sleep(0.1)
full_outputs = job.outputs()
if verbose:
print('job.outputs: %s' % str(full_outputs))
# ensure get ending to avoid race
# -1 means last response if streaming
# 0 means get text_output, ignore exception_text
# 0 means get list within text_output that looks like [[prompt], [answer]]
# 1 means get bot answer, so will have last bot answer
res_dict['response'] = md_to_text(full_outputs[-1][0][0][1], do_md_to_text=do_md_to_text)
return res_dict, client
@pytest.mark.skip(reason="For manual use against some server, no server launched")
def test_client_nochat_stream(prompt_type='human_bot'):
return run_client_nochat_gen(prompt="Tell a very long kid's story about birds.", prompt_type=prompt_type,
stream_output=True, max_new_tokens=512,
langchain_mode='Disabled')
def run_client_nochat_gen(prompt, prompt_type, stream_output, max_new_tokens, langchain_mode):
client = get_client(serialize=False)
kwargs, args = get_args(prompt, prompt_type, chat=False, stream_output=stream_output,
max_new_tokens=max_new_tokens, langchain_mode=langchain_mode)
return run_client_gen(client, prompt, args, kwargs)
def run_client_gen(client, prompt, args, kwargs, do_md_to_text=True, verbose=False):
res_dict = kwargs
res_dict['prompt'] = prompt
if not kwargs['stream_output']:
res = client.predict(str(dict(kwargs)), api_name='/submit_nochat_api')
res_dict['response'] = res[0]
print(md_to_text(res_dict['response'], do_md_to_text=do_md_to_text))
return res_dict, client
else:
job = client.submit(str(dict(kwargs)), api_name='/submit_nochat_api')
while not job.done():
outputs_list = job.communicator.job.outputs
if outputs_list:
res = job.communicator.job.outputs[-1]
res_dict = ast.literal_eval(res)
print('Stream: %s' % res_dict['response'])
time.sleep(0.1)
res_list = job.outputs()
assert len(res_list) > 0, "No response, check server"
res = res_list[-1]
res_dict = ast.literal_eval(res)
print('Final: %s' % res_dict['response'])
return res_dict, client
def md_to_text(md, do_md_to_text=True):
if not do_md_to_text:
return md
assert md is not None, "Markdown is None"
html = markdown.markdown(md)
soup = BeautifulSoup(html, features='html.parser')
return soup.get_text()
def run_client_many(prompt_type='human_bot'):
ret1, _ = test_client_chat(prompt_type=prompt_type)
ret2, _ = test_client_chat_stream(prompt_type=prompt_type)
ret3, _ = test_client_nochat_stream(prompt_type=prompt_type)
ret4, _ = test_client_basic(prompt_type=prompt_type)
ret5, _ = test_client_basic_api(prompt_type=prompt_type)
ret6, _ = test_client_basic_api_lean(prompt_type=prompt_type)
ret7, _ = test_client_basic_api_lean_morestuff(prompt_type=prompt_type)
return ret1, ret2, ret3, ret4, ret5, ret6, ret7
if __name__ == '__main__':
run_client_many()