wehub-resource-sync
115 lines
4.2 KiB
Python
115 lines
4.2 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import os
|
|
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
|
|
# os.environ['SWIFT_DEBUG'] = '1'
|
|
|
|
|
|
def infer(engine: 'InferEngine', infer_request: 'InferRequest'):
|
|
stop = [engine.template.agent_template.keyword.observation] # compat react_en
|
|
request_config = RequestConfig(max_tokens=512, temperature=0, stop=stop)
|
|
resp_list = engine.infer([infer_request], request_config)
|
|
query = infer_request.messages[0]['content']
|
|
response = resp_list[0].choices[0].message.content
|
|
print(f'query: {query}')
|
|
print(f'response: {response}')
|
|
print(f'tool_calls: {resp_list[0].choices[0].message.tool_calls}')
|
|
|
|
tool = '{"temperature": 32, "condition": "Sunny", "humidity": 50}'
|
|
print(f'tool_response: {tool}')
|
|
infer_request.messages += [{'role': 'assistant', 'content': response}, {'role': 'tool', 'content': tool}]
|
|
resp_list = engine.infer([infer_request], request_config)
|
|
response2 = resp_list[0].choices[0].message.content
|
|
print(f'response2: {response2}')
|
|
|
|
|
|
def infer_stream(engine: 'InferEngine', infer_request: 'InferRequest'):
|
|
stop = [engine.template.agent_template.keyword.observation]
|
|
request_config = RequestConfig(max_tokens=512, temperature=0, stream=True, stop=stop)
|
|
gen_list = engine.infer([infer_request], request_config)
|
|
query = infer_request.messages[0]['content']
|
|
response = ''
|
|
print(f'query: {query}\nresponse: ', end='')
|
|
for resp in gen_list[0]:
|
|
if resp is None:
|
|
continue
|
|
delta = resp.choices[0].delta.content
|
|
response += delta
|
|
print(delta, end='', flush=True)
|
|
print()
|
|
print(f'tool_calls: {resp.choices[0].delta.tool_calls}')
|
|
|
|
tool = '{"temperature": 32, "condition": "Sunny", "humidity": 50}'
|
|
print(f'tool_response: {tool}\nresponse2: ', end='')
|
|
infer_request.messages += [{'role': 'assistant', 'content': response}, {'role': 'tool', 'content': tool}]
|
|
gen_list = engine.infer([infer_request], request_config)
|
|
for resp in gen_list[0]:
|
|
if resp is None:
|
|
continue
|
|
print(resp.choices[0].delta.content, end='', flush=True)
|
|
print()
|
|
|
|
|
|
def get_infer_request():
|
|
return InferRequest(
|
|
messages=[{
|
|
'role': 'user',
|
|
'content': "How's the weather in Beijing today?"
|
|
}],
|
|
tools=[{
|
|
'name': 'get_current_weather',
|
|
'description': 'Get the current weather in a given location',
|
|
'parameters': {
|
|
'type': 'object',
|
|
'properties': {
|
|
'location': {
|
|
'type': 'string',
|
|
'description': 'The city and state, e.g. San Francisco, CA'
|
|
},
|
|
'unit': {
|
|
'type': 'string',
|
|
'enum': ['celsius', 'fahrenheit']
|
|
}
|
|
},
|
|
'required': ['location']
|
|
}
|
|
}])
|
|
|
|
|
|
def infer_continue_generate(engine):
|
|
# Continue generating after the assistant message.
|
|
infer_request = InferRequest(messages=[{
|
|
'role': 'user',
|
|
'content': 'How is the weather today?'
|
|
}, {
|
|
'role': 'assistant',
|
|
'content': 'It is sunny today, '
|
|
}])
|
|
request_config = RequestConfig(max_tokens=512, temperature=0)
|
|
resp_list = engine.infer([infer_request], request_config)
|
|
response = resp_list[0].choices[0].message.content
|
|
print(f'response: {response}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
from swift.agent_template import agent_template_map
|
|
from swift.infer_engine import InferEngine, InferRequest, RequestConfig, TransformersEngine
|
|
model = 'Qwen/Qwen2.5-1.5B-Instruct'
|
|
infer_backend = 'transformers'
|
|
|
|
if infer_backend == 'transformers':
|
|
engine = TransformersEngine(model, max_batch_size=64)
|
|
elif infer_backend == 'vllm':
|
|
from swift.infer_engine import VllmEngine
|
|
engine = VllmEngine(model, max_model_len=8192)
|
|
elif infer_backend == 'lmdeploy':
|
|
from swift.infer_engine import LmdeployEngine
|
|
engine = LmdeployEngine(model)
|
|
|
|
# engine.template._agent_template = 'hermes' # react_en/qwen_en/qwen_en_parallel
|
|
|
|
infer(engine, get_infer_request())
|
|
infer_stream(engine, get_infer_request())
|
|
|
|
# infer_continue_generate(engine)
|