Function _test_client

tests/deploy/test_logprobs.py:1–63

(port: int, print_logprobs: bool = False, test_vlm: bool = False)

def _test_client(port: int, print_logprobs: bool = False, test_vlm: bool = False):
    import aiohttp
    import time
    from pprint import pprint

    from swift.infer_engine import InferClient, InferRequest, RequestConfig

    infer_client = InferClient(port=port)

    while True:
        try:
            models = infer_client.models
            print(f'models: {models}')
        except aiohttp.ClientConnectorError:
            time.sleep(5)
            continue
        break

    if test_vlm:
        query = '这是什么'
        # http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/cat.png
        messages = [{
            'role':
            'user',
            'content': [
                {
                    'type': 'text',
                    'text': '这是什么'
                },
                {
                    'type': 'image_url',
                    'image_url': {
                        'url': 'cat.png'
                    }
                },
            ]
        }]
    else:
        query = '123*234=?'
        messages = [{'role': 'user', 'content': query}]

    infer_request = InferRequest(messages=messages)
    request_config = RequestConfig(seed=42, max_tokens=256, temperature=0.8, logprobs=True, top_logprobs=5)

    resp = infer_client.infer([infer_request], request_config=request_config)[0]
    response = resp.choices[0].message.content
    print(f'query: {query}')
    print(f'response: {response}')
    if print_logprobs:
        pprint(resp.choices[0].logprobs)

    request_config = RequestConfig(
        stream=True, seed=42, max_tokens=256, temperature=0.8, top_k=20, top_p=0.8, logprobs=True, top_logprobs=5)
    gen_list = infer_client.infer([infer_request], request_config=request_config)
    print(f'query: {query}')
    print('response: ', end='')
    for chunk in gen_list[0]:
        if chunk is None:
            continue
        print(chunk.choices[0].delta.content, end='', flush=True)