Spaces:
Sleeping
Sleeping
| import asyncio | |
| import copy | |
| import logging | |
| from dataclasses import asdict | |
| from typing import List, Optional, Union | |
| import aiohttp | |
| from lagent.llms.base_llm import AsyncLLMMixin, BaseLLM | |
| from lagent.schema import ModelStatusCode | |
| from lagent.utils.util import filter_suffix | |
| class TritonClient(BaseLLM): | |
| """TritonClient is a wrapper of TritonClient for LLM. | |
| Args: | |
| tritonserver_addr (str): the address in format "ip:port" of | |
| triton inference server | |
| model_name (str): the name of the model | |
| session_len (int): the context size | |
| max_tokens (int): the expected generated token numbers | |
| """ | |
| def __init__(self, | |
| tritonserver_addr: str, | |
| model_name: str, | |
| session_len: int = 32768, | |
| log_level: str = 'WARNING', | |
| **kwargs): | |
| super().__init__(path=None, **kwargs) | |
| try: | |
| from lmdeploy.serve.turbomind.chatbot import Chatbot, StatusCode | |
| except Exception as e: | |
| logging.error(f'{e}') | |
| raise RuntimeError('DO NOT use turbomind.chatbot since it has ' | |
| 'been removed by lmdeploy since v0.5.2') | |
| self.state_map = { | |
| StatusCode.TRITON_STREAM_END: ModelStatusCode.END, | |
| StatusCode.TRITON_SERVER_ERR: ModelStatusCode.SERVER_ERR, | |
| StatusCode.TRITON_SESSION_CLOSED: ModelStatusCode.SESSION_CLOSED, | |
| StatusCode.TRITON_STREAM_ING: ModelStatusCode.STREAM_ING, | |
| StatusCode.TRITON_SESSION_OUT_OF_LIMIT: | |
| ModelStatusCode.SESSION_OUT_OF_LIMIT, | |
| StatusCode.TRITON_SESSION_INVALID_ARG: | |
| ModelStatusCode.SESSION_INVALID_ARG, | |
| StatusCode.TRITON_SESSION_READY: ModelStatusCode.SESSION_READY | |
| } | |
| self.chatbot = Chatbot( | |
| tritonserver_addr=tritonserver_addr, | |
| model_name=model_name, | |
| session_len=session_len, | |
| log_level=log_level, | |
| **kwargs) | |
| def generate(self, | |
| inputs: Union[str, List[str]], | |
| session_id: int = 2967, | |
| request_id: str = '', | |
| sequence_start: bool = True, | |
| sequence_end: bool = True, | |
| skip_special_tokens: bool = False, | |
| **kwargs): | |
| """Start a new round conversation of a session. Return the chat | |
| completions in non-stream mode. | |
| Args: | |
| inputs (str, List[str]): user's prompt(s) in this round | |
| session_id (int): the identical id of a session | |
| request_id (str): the identical id of this round conversation | |
| sequence_start (bool): start flag of a session | |
| sequence_end (bool): end flag of a session | |
| skip_special_tokens (bool): Whether or not to remove special tokens | |
| in the decoding. Default to be False. | |
| Returns: | |
| (a list of/batched) text/chat completion | |
| """ | |
| from lmdeploy.serve.turbomind.chatbot import Session, get_logger | |
| if isinstance(inputs, str): | |
| inputs = [inputs] | |
| prompt = inputs | |
| assert isinstance(session_id, int), \ | |
| f'INT session id is required, but got {type(session_id)}' | |
| self.chatbot.cfg = self._update_gen_params(**kwargs) | |
| max_new_tokens = self.chatbot.cfg.max_new_tokens | |
| logger = get_logger('service.ft', log_level=self.chatbot.log_level) | |
| logger.info(f'session {session_id}, request_id {request_id}, ' | |
| f'max_out_len {max_new_tokens}') | |
| if self.chatbot._session is None: | |
| sequence_start = True | |
| self.chatbot._session = Session(session_id=session_id) | |
| elif self.chatbot._session.status == 0: | |
| logger.error(f'session {session_id} has been ended. Please set ' | |
| f'`sequence_start` be True if you want to restart it') | |
| return '' | |
| self.chatbot._session.status = 1 | |
| self.chatbot._session.request_id = request_id | |
| self.chatbot._session.response = '' | |
| status, res, _ = None, '', 0 | |
| for status, res, _ in self.chatbot._stream_infer( | |
| self.chatbot._session, | |
| prompt, | |
| max_new_tokens, | |
| sequence_start, | |
| sequence_end, | |
| skip_special_tokens=skip_special_tokens): | |
| status = self.state_map.get(status) | |
| if status < ModelStatusCode.END: | |
| return '' | |
| elif status == ModelStatusCode.END: | |
| self.chatbot._session.histories = ( | |
| self.chatbot._session.histories + | |
| self.chatbot._session.prompt + | |
| self.chatbot._session.response) | |
| # remove stop_words | |
| res = filter_suffix(res, self.gen_params.get('stop_words')) | |
| return res | |
| def stream_chat(self, | |
| inputs: List[dict], | |
| session_id: int = 2967, | |
| request_id: str = '', | |
| sequence_start: bool = True, | |
| sequence_end: bool = True, | |
| skip_special_tokens: bool = False, | |
| **kwargs): | |
| """Start a new round conversation of a session. Return the chat | |
| completions in stream mode. | |
| Args: | |
| session_id (int): the identical id of a session | |
| inputs (List[dict]): user's inputs in this round conversation | |
| request_id (str): the identical id of this round conversation | |
| sequence_start (bool): start flag of a session | |
| sequence_end (bool): end flag of a session | |
| skip_special_tokens (bool): Whether or not to remove special tokens | |
| in the decoding. Default to be False. | |
| Returns: | |
| tuple(Status, str, int): status, text/chat completion, | |
| generated token number | |
| """ | |
| from lmdeploy.serve.turbomind.chatbot import Session, get_logger | |
| assert isinstance(session_id, int), \ | |
| f'INT session id is required, but got {type(session_id)}' | |
| self.chatbot.cfg = self._update_gen_params(**kwargs) | |
| max_new_tokens = self.chatbot.cfg.max_new_tokens | |
| logger = get_logger('service.ft', log_level=self.chatbot.log_level) | |
| logger.info(f'session {session_id}, request_id {request_id}, ' | |
| f'max_out_len {max_new_tokens}') | |
| if self.chatbot._session is None: | |
| sequence_start = True | |
| self.chatbot._session = Session(session_id=session_id) | |
| elif self.chatbot._session.status == 0: | |
| logger.error(f'session {session_id} has been ended. Please set ' | |
| f'`sequence_start` be True if you want to restart it') | |
| return ModelStatusCode.SESSION_CLOSED, '', 0 | |
| self.chatbot._session.status = 1 | |
| self.chatbot._session.request_id = request_id | |
| self.chatbot._session.response = '' | |
| prompt = self.template_parser(inputs) | |
| status, res, _ = None, '', 0 | |
| for status, res, _ in self.chatbot._stream_infer( | |
| self.chatbot._session, | |
| prompt, | |
| max_new_tokens, | |
| sequence_start, | |
| sequence_end, | |
| skip_special_tokens=skip_special_tokens): | |
| status = self.state_map.get(status) | |
| # The stop symbol also appears in the output of the last STREAM_ING state. | |
| res = filter_suffix(res, self.gen_params.get('stop_words')) | |
| if status < ModelStatusCode.END: | |
| return status, res, _ | |
| elif status == ModelStatusCode.END: # remove stop_words | |
| self.chatbot._session.histories = ( | |
| self.chatbot._session.histories + | |
| self.chatbot._session.prompt + | |
| self.chatbot._session.response) | |
| yield status, res, _ | |
| break | |
| else: | |
| yield status, res, _ | |
| def _update_gen_params(self, **kwargs): | |
| import mmengine | |
| new_gen_params = self.update_gen_params(**kwargs) | |
| self.gen_params['stop_words'] = new_gen_params.pop('stop_words') | |
| stop_words = self.chatbot._stop_words( | |
| self.gen_params.get('stop_words')) | |
| cfg = mmengine.Config( | |
| dict( | |
| session_len=self.chatbot.model.session_len, | |
| stop_words=stop_words, | |
| bad_words=self.chatbot.cfg.bad_words, | |
| **new_gen_params)) | |
| return cfg | |
| class LMDeployPipeline(BaseLLM): | |
| """ | |
| Args: | |
| path (str): The path to the model. | |
| It could be one of the following options: | |
| - i) A local directory path of a turbomind model which is | |
| converted by `lmdeploy convert` command or download | |
| from ii) and iii). | |
| - ii) The model_id of a lmdeploy-quantized model hosted | |
| inside a model repo on huggingface.co, such as | |
| "InternLM/internlm-chat-20b-4bit", | |
| "lmdeploy/llama2-chat-70b-4bit", etc. | |
| - iii) The model_id of a model hosted inside a model repo | |
| on huggingface.co, such as "internlm/internlm-chat-7b", | |
| "Qwen/Qwen-7B-Chat ", "baichuan-inc/Baichuan2-7B-Chat" | |
| and so on. | |
| model_name (str): needed when model_path is a pytorch model on | |
| huggingface.co, such as "internlm-chat-7b", | |
| "Qwen-7B-Chat ", "Baichuan2-7B-Chat" and so on. | |
| tp (int): tensor parallel | |
| pipeline_cfg (dict): config of pipeline | |
| """ | |
| def __init__(self, | |
| path: str, | |
| model_name: Optional[str] = None, | |
| tp: int = 1, | |
| pipeline_cfg=dict(), | |
| **kwargs): | |
| import lmdeploy | |
| from lmdeploy import ChatTemplateConfig, TurbomindEngineConfig, pipeline, version_info | |
| self.str_version = lmdeploy.__version__ | |
| self.version = version_info | |
| self.do_sample = kwargs.pop('do_sample', None) | |
| if self.do_sample is not None and self.version < (0, 6, 0): | |
| raise RuntimeError( | |
| '`do_sample` parameter is not supported by lmdeploy until ' | |
| f'v0.6.0, but currently using lmdeloy {self.str_version}') | |
| super().__init__(path=path, **kwargs) | |
| backend_config = copy.deepcopy(pipeline_cfg) | |
| backend_config.update(tp=tp) | |
| backend_config = { | |
| k: v | |
| for k, v in backend_config.items() | |
| if hasattr(TurbomindEngineConfig, k) | |
| } | |
| backend_config = TurbomindEngineConfig(**backend_config) | |
| chat_template_config = ChatTemplateConfig( | |
| model_name=model_name) if model_name else None | |
| self.model = pipeline( | |
| model_path=self.path, | |
| backend_config=backend_config, | |
| chat_template_config=chat_template_config, | |
| log_level='WARNING') | |
| def generate(self, | |
| inputs: Union[str, List[str]], | |
| do_preprocess: bool = None, | |
| skip_special_tokens: bool = False, | |
| return_dict: bool = False, | |
| **kwargs): | |
| """Return the chat completions in non-stream mode. | |
| Args: | |
| inputs (Union[str, List[str]]): input texts to be completed. | |
| do_preprocess (bool): whether pre-process the messages. Default to | |
| True, which means chat_template will be applied. | |
| skip_special_tokens (bool): Whether or not to remove special tokens | |
| in the decoding. Default to be False. | |
| Returns: | |
| (a list of/batched) text/chat completion | |
| """ | |
| from lmdeploy.messages import GenerationConfig | |
| batched = True | |
| if isinstance(inputs, str): | |
| inputs = [inputs] | |
| batched = False | |
| prompt = inputs | |
| do_sample = kwargs.pop('do_sample', None) | |
| gen_params = self.update_gen_params(**kwargs) | |
| if do_sample is None: | |
| do_sample = self.do_sample | |
| if do_sample is not None and self.version < (0, 6, 0): | |
| raise RuntimeError( | |
| '`do_sample` parameter is not supported by lmdeploy until ' | |
| f'v0.6.0, but currently using lmdeloy {self.str_version}') | |
| if self.version >= (0, 6, 0): | |
| if do_sample is None: | |
| do_sample = gen_params['top_k'] > 1 or gen_params[ | |
| 'temperature'] > 0 | |
| gen_params.update(do_sample=do_sample) | |
| gen_config = GenerationConfig( | |
| skip_special_tokens=skip_special_tokens, **gen_params) | |
| response = self.model.batch_infer( | |
| prompt, gen_config=gen_config, do_preprocess=do_preprocess) | |
| texts = [resp.text for resp in response] | |
| # remove stop_words | |
| texts = filter_suffix(texts, self.gen_params.get('stop_words')) | |
| for resp, text in zip(response, texts): | |
| resp.text = text | |
| if batched: | |
| return [asdict(resp) | |
| for resp in response] if return_dict else texts | |
| return asdict(response[0]) if return_dict else texts[0] | |
| class LMDeployServer(BaseLLM): | |
| """ | |
| Args: | |
| path (str): The path to the model. | |
| It could be one of the following options: | |
| - i) A local directory path of a turbomind model which is | |
| converted by `lmdeploy convert` command or download from | |
| ii) and iii). | |
| - ii) The model_id of a lmdeploy-quantized model hosted | |
| inside a model repo on huggingface.co, such as | |
| "InternLM/internlm-chat-20b-4bit", | |
| "lmdeploy/llama2-chat-70b-4bit", etc. | |
| - iii) The model_id of a model hosted inside a model repo | |
| on huggingface.co, such as "internlm/internlm-chat-7b", | |
| "Qwen/Qwen-7B-Chat ", "baichuan-inc/Baichuan2-7B-Chat" | |
| and so on. | |
| model_name (str): needed when model_path is a pytorch model on | |
| huggingface.co, such as "internlm-chat-7b", | |
| "Qwen-7B-Chat ", "Baichuan2-7B-Chat" and so on. | |
| server_name (str): host ip for serving | |
| server_port (int): server port | |
| tp (int): tensor parallel | |
| log_level (str): set log level whose value among | |
| [CRITICAL, ERROR, WARNING, INFO, DEBUG] | |
| """ | |
| def __init__(self, | |
| path: str, | |
| model_name: Optional[str] = None, | |
| server_name: str = '0.0.0.0', | |
| server_port: int = 23333, | |
| tp: int = 1, | |
| log_level: str = 'WARNING', | |
| serve_cfg=dict(), | |
| **kwargs): | |
| super().__init__(path=path, **kwargs) | |
| self.model_name = model_name | |
| # TODO get_logger issue in multi processing | |
| import lmdeploy | |
| self.client = lmdeploy.serve( | |
| model_path=self.path, | |
| model_name=model_name, | |
| server_name=server_name, | |
| server_port=server_port, | |
| tp=tp, | |
| log_level=log_level, | |
| **serve_cfg) | |
| def generate(self, | |
| inputs: Union[str, List[str]], | |
| session_id: int = 2967, | |
| sequence_start: bool = True, | |
| sequence_end: bool = True, | |
| ignore_eos: bool = False, | |
| skip_special_tokens: Optional[bool] = False, | |
| timeout: int = 30, | |
| **kwargs) -> List[str]: | |
| """Start a new round conversation of a session. Return the chat | |
| completions in non-stream mode. | |
| Args: | |
| inputs (str, List[str]): user's prompt(s) in this round | |
| session_id (int): the identical id of a session | |
| sequence_start (bool): start flag of a session | |
| sequence_end (bool): end flag of a session | |
| ignore_eos (bool): indicator for ignoring eos | |
| skip_special_tokens (bool): Whether or not to remove special tokens | |
| in the decoding. Default to be False. | |
| timeout (int): max time to wait for response | |
| Returns: | |
| (a list of/batched) text/chat completion | |
| """ | |
| batched = True | |
| if isinstance(inputs, str): | |
| inputs = [inputs] | |
| batched = False | |
| gen_params = self.update_gen_params(**kwargs) | |
| max_new_tokens = gen_params.pop('max_new_tokens') | |
| gen_params.update(max_tokens=max_new_tokens) | |
| resp = [''] * len(inputs) | |
| for text in self.client.completions_v1( | |
| self.model_name, | |
| inputs, | |
| session_id=session_id, | |
| sequence_start=sequence_start, | |
| sequence_end=sequence_end, | |
| stream=False, | |
| ignore_eos=ignore_eos, | |
| skip_special_tokens=skip_special_tokens, | |
| timeout=timeout, | |
| **gen_params): | |
| resp = [ | |
| resp[i] + item['text'] | |
| for i, item in enumerate(text['choices']) | |
| ] | |
| # remove stop_words | |
| resp = filter_suffix(resp, self.gen_params.get('stop_words')) | |
| if not batched: | |
| return resp[0] | |
| return resp | |
| def stream_chat(self, | |
| inputs: List[dict], | |
| session_id=0, | |
| sequence_start: bool = True, | |
| sequence_end: bool = True, | |
| stream: bool = True, | |
| ignore_eos: bool = False, | |
| skip_special_tokens: Optional[bool] = False, | |
| timeout: int = 30, | |
| **kwargs): | |
| """Start a new round conversation of a session. Return the chat | |
| completions in stream mode. | |
| Args: | |
| session_id (int): the identical id of a session | |
| inputs (List[dict]): user's inputs in this round conversation | |
| sequence_start (bool): start flag of a session | |
| sequence_end (bool): end flag of a session | |
| stream (bool): return in a streaming format if enabled | |
| ignore_eos (bool): indicator for ignoring eos | |
| skip_special_tokens (bool): Whether or not to remove special tokens | |
| in the decoding. Default to be False. | |
| timeout (int): max time to wait for response | |
| Returns: | |
| tuple(Status, str, int): status, text/chat completion, | |
| generated token number | |
| """ | |
| gen_params = self.update_gen_params(**kwargs) | |
| max_new_tokens = gen_params.pop('max_new_tokens') | |
| gen_params.update(max_tokens=max_new_tokens) | |
| prompt = self.template_parser(inputs) | |
| resp = '' | |
| finished = False | |
| stop_words = self.gen_params.get('stop_words') | |
| for text in self.client.completions_v1( | |
| self.model_name, | |
| prompt, | |
| session_id=session_id, | |
| sequence_start=sequence_start, | |
| sequence_end=sequence_end, | |
| stream=stream, | |
| ignore_eos=ignore_eos, | |
| skip_special_tokens=skip_special_tokens, | |
| timeout=timeout, | |
| **gen_params): | |
| resp += text['choices'][0]['text'] | |
| if not resp: | |
| continue | |
| # remove stop_words | |
| for sw in stop_words: | |
| if sw in resp: | |
| resp = filter_suffix(resp, stop_words) | |
| finished = True | |
| break | |
| yield ModelStatusCode.STREAM_ING, resp, None | |
| if finished: | |
| break | |
| yield ModelStatusCode.END, resp, None | |
| class LMDeployClient(LMDeployServer): | |
| """ | |
| Args: | |
| url (str): communicating address 'http://<ip>:<port>' of | |
| api_server | |
| model_name (str): needed when model_path is a pytorch model on | |
| huggingface.co, such as "internlm-chat-7b", | |
| "Qwen-7B-Chat ", "Baichuan2-7B-Chat" and so on. | |
| """ | |
| def __init__(self, url: str, model_name: str, **kwargs): | |
| BaseLLM.__init__(self, path=url, **kwargs) | |
| from lmdeploy.serve.openai.api_client import APIClient | |
| self.client = APIClient(url) | |
| self.model_name = model_name | |
| class AsyncLMDeployPipeline(AsyncLLMMixin, LMDeployPipeline): | |
| """ | |
| Args: | |
| path (str): The path to the model. | |
| It could be one of the following options: | |
| - i) A local directory path of a turbomind model which is | |
| converted by `lmdeploy convert` command or download | |
| from ii) and iii). | |
| - ii) The model_id of a lmdeploy-quantized model hosted | |
| inside a model repo on huggingface.co, such as | |
| "InternLM/internlm-chat-20b-4bit", | |
| "lmdeploy/llama2-chat-70b-4bit", etc. | |
| - iii) The model_id of a model hosted inside a model repo | |
| on huggingface.co, such as "internlm/internlm-chat-7b", | |
| "Qwen/Qwen-7B-Chat ", "baichuan-inc/Baichuan2-7B-Chat" | |
| and so on. | |
| model_name (str): needed when model_path is a pytorch model on | |
| huggingface.co, such as "internlm-chat-7b", | |
| "Qwen-7B-Chat ", "Baichuan2-7B-Chat" and so on. | |
| tp (int): tensor parallel | |
| pipeline_cfg (dict): config of pipeline | |
| """ | |
| async def generate(self, | |
| inputs: Union[str, List[str]], | |
| session_ids: Union[int, List[int]] = None, | |
| do_preprocess: bool = None, | |
| skip_special_tokens: bool = False, | |
| return_dict: bool = False, | |
| **kwargs): | |
| """Return the chat completions in non-stream mode. | |
| Args: | |
| inputs (Union[str, List[str]]): input texts to be completed. | |
| do_preprocess (bool): whether pre-process the messages. Default to | |
| True, which means chat_template will be applied. | |
| skip_special_tokens (bool): Whether or not to remove special tokens | |
| in the decoding. Default to be False. | |
| Returns: | |
| (a list of/batched) text/chat completion | |
| """ | |
| from lmdeploy.messages import GenerationConfig, Response | |
| batched = True | |
| if isinstance(inputs, str): | |
| inputs = [inputs] | |
| batched = False | |
| if session_ids is None: | |
| session_ids = list(range(len(inputs))) | |
| elif isinstance(session_ids, (int, str)): | |
| session_ids = [session_ids] | |
| assert len(inputs) == len(session_ids) | |
| prompt = inputs | |
| gen_params = self.update_gen_params(**kwargs) | |
| gen_config = GenerationConfig( | |
| skip_special_tokens=skip_special_tokens, **gen_params) | |
| async def _inner_generate(uid, text): | |
| resp = Response('', 0, 0, uid) | |
| async for out in self.model.generate( | |
| text, | |
| uid, | |
| gen_config, | |
| stream_response=True, | |
| sequence_start=True, | |
| sequence_end=True, | |
| do_preprocess=do_preprocess, | |
| **kwargs): | |
| resp.text += out.response | |
| resp.generate_token_len = out.generate_token_len | |
| resp.input_token_len = out.input_token_len | |
| resp.finish_reason = out.finish_reason | |
| if out.token_ids: | |
| resp.token_ids.extend(out.token_ids) | |
| if out.logprobs: | |
| if resp.logprobs is None: | |
| resp.logprobs = [] | |
| resp.logprobs.extend(out.logprobs) | |
| return resp | |
| response = await asyncio.gather(*[ | |
| _inner_generate(sid, inp) for sid, inp in zip(session_ids, prompt) | |
| ]) | |
| texts = [resp.text for resp in response] | |
| # remove stop_words | |
| texts = filter_suffix(texts, self.gen_params.get('stop_words')) | |
| for resp, text in zip(response, texts): | |
| resp.text = text | |
| if batched: | |
| return [asdict(resp) | |
| for resp in response] if return_dict else texts | |
| return asdict(response[0]) if return_dict else texts[0] | |
| class AsyncLMDeployServer(AsyncLLMMixin, LMDeployServer): | |
| """ | |
| Args: | |
| path (str): The path to the model. | |
| It could be one of the following options: | |
| - i) A local directory path of a turbomind model which is | |
| converted by `lmdeploy convert` command or download from | |
| ii) and iii). | |
| - ii) The model_id of a lmdeploy-quantized model hosted | |
| inside a model repo on huggingface.co, such as | |
| "InternLM/internlm-chat-20b-4bit", | |
| "lmdeploy/llama2-chat-70b-4bit", etc. | |
| - iii) The model_id of a model hosted inside a model repo | |
| on huggingface.co, such as "internlm/internlm-chat-7b", | |
| "Qwen/Qwen-7B-Chat ", "baichuan-inc/Baichuan2-7B-Chat" | |
| and so on. | |
| model_name (str): needed when model_path is a pytorch model on | |
| huggingface.co, such as "internlm-chat-7b", | |
| "Qwen-7B-Chat ", "Baichuan2-7B-Chat" and so on. | |
| server_name (str): host ip for serving | |
| server_port (int): server port | |
| tp (int): tensor parallel | |
| log_level (str): set log level whose value among | |
| [CRITICAL, ERROR, WARNING, INFO, DEBUG] | |
| """ | |
| async def generate( | |
| self, | |
| inputs: Union[str, List[str]], | |
| session_ids: Union[int, List[int]] = None, | |
| sequence_start: bool = True, | |
| sequence_end: bool = True, | |
| ignore_eos: bool = False, | |
| skip_special_tokens: Optional[bool] = False, | |
| timeout: int = 30, | |
| **kwargs, | |
| ): | |
| """Start a new round conversation of a session. Return the chat | |
| completions in non-stream mode. | |
| Args: | |
| inputs (str, List[str]): user's prompt(s) in this round | |
| session_ids (int, List[int]): session id(s) | |
| sequence_start (bool): start flag of a session | |
| sequence_end (bool): end flag of a session | |
| ignore_eos (bool): indicator for ignoring eos | |
| skip_special_tokens (bool): Whether or not to remove special tokens | |
| in the decoding. Default to be False. | |
| timeout (int): max time to wait for response | |
| Returns: | |
| (a list of/batched) text/chat completion | |
| """ | |
| from lmdeploy.serve.openai.api_client import json_loads | |
| batched = True | |
| if isinstance(inputs, str): | |
| inputs = [inputs] | |
| batched = False | |
| gen_params = self.update_gen_params(**kwargs) | |
| max_new_tokens = gen_params.pop('max_new_tokens') | |
| gen_params.update(max_tokens=max_new_tokens) | |
| responses = [''] * len(inputs) | |
| pload = dict( | |
| model=self.model_name, | |
| prompt=inputs, | |
| sequence_start=sequence_start, | |
| sequence_end=sequence_end, | |
| stream=False, | |
| ignore_eos=ignore_eos, | |
| skip_special_tokens=skip_special_tokens, | |
| timeout=timeout, | |
| **gen_params) | |
| async with aiohttp.ClientSession( | |
| timeout=aiohttp.ClientTimeout(3 * 3600)) as session: | |
| async with session.post( | |
| self.client.completions_v1_url, | |
| headers=self.client.headers, | |
| json=pload) as resp: | |
| async for chunk in resp.content: | |
| if chunk: | |
| decoded = chunk.decode('utf-8') | |
| output = json_loads(decoded) | |
| responses = [ | |
| response + item['text'] for response, item in zip( | |
| responses, output['choices']) | |
| ] | |
| # remove stop_words | |
| responses = filter_suffix(responses, self.gen_params.get('stop_words')) | |
| if not batched: | |
| return responses[0] | |
| return responses | |
| async def stream_chat( | |
| self, | |
| inputs: List[dict], | |
| session_id: int = None, | |
| sequence_start: bool = True, | |
| sequence_end: bool = True, | |
| stream: bool = True, | |
| ignore_eos: bool = False, | |
| skip_special_tokens: Optional[bool] = False, | |
| timeout: int = 30, | |
| **kwargs, | |
| ): | |
| """Start a new round conversation of a session. Return the chat | |
| completions in stream mode. | |
| Args: | |
| inputs (List[dict]): user's inputs in this round conversation | |
| session_id (int): session id | |
| sequence_start (bool): start flag of a session | |
| sequence_end (bool): end flag of a session | |
| stream (bool): return in a streaming format if enabled | |
| ignore_eos (bool): indicator for ignoring eos | |
| skip_special_tokens (bool): Whether or not to remove special tokens | |
| in the decoding. Default to be False. | |
| timeout (int): max time to wait for response | |
| Returns: | |
| tuple(Status, str, int): status, text/chat completion, | |
| generated token number | |
| """ | |
| from lmdeploy.serve.openai.api_client import json_loads | |
| gen_params = self.update_gen_params(**kwargs) | |
| max_new_tokens = gen_params.pop('max_new_tokens') | |
| gen_params.update(max_tokens=max_new_tokens) | |
| prompt = self.template_parser(inputs) | |
| response = '' | |
| finished = False | |
| stop_words = self.gen_params.get('stop_words') | |
| pload = dict( | |
| model=self.model_name, | |
| prompt=prompt, | |
| sequence_start=sequence_start, | |
| sequence_end=sequence_end, | |
| stream=stream, | |
| ignore_eos=ignore_eos, | |
| skip_special_tokens=skip_special_tokens, | |
| timeout=timeout, | |
| **gen_params) | |
| async with aiohttp.ClientSession( | |
| timeout=aiohttp.ClientTimeout(3 * 3600)) as session: | |
| async with session.post( | |
| self.client.completions_v1_url, | |
| headers=self.client.headers, | |
| json=pload) as resp: | |
| async for chunk in resp.content: | |
| if chunk: | |
| decoded = chunk.decode('utf-8') | |
| if not decoded.strip() or decoded.rstrip( | |
| ) == 'data: [DONE]': | |
| continue | |
| if decoded[:6] == 'data: ': | |
| decoded = decoded[6:] | |
| output = json_loads(decoded) | |
| response += output['choices'][0]['text'] | |
| if not response: | |
| continue | |
| # remove stop_words | |
| for sw in stop_words: | |
| if sw in response: | |
| response = filter_suffix(response, stop_words) | |
| finished = True | |
| break | |
| yield ModelStatusCode.STREAM_ING, response, None | |
| if finished: | |
| break | |
| yield ModelStatusCode.END, response, None | |
| class AsyncLMDeployClient(AsyncLMDeployServer): | |
| """ | |
| Args: | |
| url (str): communicating address 'http://<ip>:<port>' of | |
| api_server | |
| model_name (str): needed when model_path is a pytorch model on | |
| huggingface.co, such as "internlm-chat-7b", | |
| "Qwen-7B-Chat ", "Baichuan2-7B-Chat" and so on. | |
| """ | |
| def __init__(self, url: str, model_name: str, **kwargs): | |
| BaseLLM.__init__(self, path=url, **kwargs) | |
| from lmdeploy.serve.openai.api_client import APIClient | |
| self.client = APIClient(url) | |
| self.model_name = model_name | |