Function main

lmdeploy/cli/chat.py:72–107

(model_path, backend, trust_remote_code=False, **kwargs)

def main(model_path, backend, trust_remote_code=False, **kwargs):
    if backend != 'pytorch':
        # set auto backend mode
        backend = autoget_backend(model_path, trust_remote_code=trust_remote_code)
    quit = False
    with build_pipe(model_path, backend, trust_remote_code=trust_remote_code, **kwargs) as pipe:
        gen_config = build_gen_config(**kwargs)
        adapter_name = get_adapter_name(**kwargs)
        while not quit:
            with closing(pipe.session()) as sess:
                while True:
                    try:
                        prompt = input_prompt()
                    except KeyboardInterrupt:
                        quit = True
                        break
                    if prompt == 'end':
                        sess.close()
                        break
                    if prompt == 'exit':
                        quit = True
                        break
                    if prompt.strip() == '':
                        continue
                    resps = pipe.chat(prompt,
                                      session=sess,
                                      gen_config=gen_config,
                                      adapter_name=adapter_name,
                                      stream_response=True)
                    try:
                        for resp in resps:
                            print(resp.text, end='', flush=True)
                    except KeyboardInterrupt:
                        sess.abort()
        else:
            print('exiting...')


if __name__ == '__main__':