| | from copy import copy |
| | from typing import Dict, List, Optional, Tuple, Union |
| |
|
| |
|
| | class LMTemplateParser: |
| | """Intermidate prompt template parser, specifically for language models. |
| | |
| | Args: |
| | meta_template (list of dict, optional): The meta template for the |
| | model. |
| | """ |
| |
|
| | def __init__(self, meta_template: Optional[List[Dict]] = None): |
| | self.meta_template = meta_template |
| | if meta_template: |
| | assert isinstance(meta_template, list) |
| | self.roles: Dict[str, dict] = dict() |
| | for item in meta_template: |
| | assert isinstance(item, dict) |
| | assert item['role'] not in self.roles, \ |
| | 'role in meta prompt must be unique!' |
| | self.roles[item['role']] = item.copy() |
| |
|
| | def __call__(self, dialog) -> str: |
| | """Parse a prompt template, and wrap it with meta template if |
| | applicable. |
| | |
| | Args: |
| | dialog (List[str or PromptList]): A prompt |
| | template (potentially before being wrapped by meta template). |
| | |
| | Returns: |
| | str: The final string. |
| | """ |
| | assert isinstance(dialog, (str, list)) |
| | if isinstance(dialog, str): |
| | return dialog |
| | if self.meta_template: |
| |
|
| | prompt = '' |
| | for index, item in enumerate(dialog): |
| | if isinstance(item, str): |
| | prompt += item |
| | else: |
| | new_str = self._prompt2str(item, index == len(dialog) - 1) |
| | prompt += new_str |
| | else: |
| | |
| | prompt = '' |
| | last_sep = '' |
| | for item in dialog: |
| | if isinstance(item, str): |
| | if item: |
| | prompt += last_sep + item |
| | elif item.get('content', ''): |
| | prompt += last_sep + item.get('prompt', '') |
| | last_sep = '\n' |
| | return prompt |
| |
|
| | def _format_begin(self, role_cfg, message): |
| | name = message.get('name', None) |
| | if name is not None: |
| | begin = role_cfg['begin'].get('with_name', '') |
| | if name in role_cfg['begin'].get('name', {}): |
| | begin = begin.format(name=role_cfg['begin']['name'][name]) |
| | else: |
| | begin = begin.format(name=name) |
| | else: |
| | if isinstance(role_cfg.get('begin', ''), str): |
| | begin = role_cfg.get('begin', '') |
| | elif isinstance(role_cfg['begin'], dict): |
| | begin = role_cfg['begin'].get('without_name', '') |
| | return begin |
| |
|
| | def _prompt2str(self, |
| | prompt: Union[str, Dict], |
| | last: bool = False) -> Tuple[str, bool]: |
| | if isinstance(prompt, str): |
| | return prompt |
| | merged_prompt = self.roles.get(prompt['role']) |
| |
|
| | if merged_prompt.get('fallback_role'): |
| | merged_prompt = self.roles.get(merged_prompt['fallback_role']) |
| | begin = self._format_begin(merged_prompt, prompt) |
| | res = begin |
| | if last and merged_prompt.get('generate', False): |
| | res += prompt.get('content', '') |
| | return res |
| | res += prompt.get('content', '') + merged_prompt.get('end', '') |
| | if last and merged_prompt['role'] != 'assistant': |
| | res += self._format_begin(self.roles['assistant'], {}) |
| | return res |
| | return res |
| |
|
| |
|
| | class BaseLLM: |
| | """Base class for model wrapper. |
| | |
| | Args: |
| | path (str): The path to the model. |
| | max_new_tokens (int): Maximum length of output expected to be generated by the model. Defaults |
| | to 512. |
| | tokenizer_only (bool): If True, only the tokenizer will be initialized. |
| | Defaults to False. |
| | meta_template (list of dict, optional): The model's meta prompt |
| | template if needed, in case the requirement of injecting or |
| | wrapping of any meta instructions. |
| | """ |
| |
|
| | def __init__(self, |
| | path: str, |
| | tokenizer_only: bool = False, |
| | template_parser: 'LMTemplateParser' = LMTemplateParser, |
| | meta_template: Optional[List[Dict]] = None, |
| | *, |
| | max_new_tokens: int = 512, |
| | top_p: float = 0.8, |
| | top_k: float = 40, |
| | temperature: float = 0.8, |
| | repetition_penalty: float = 1.0, |
| | stop_words: Union[List[str], str] = None): |
| | self.path = path |
| | self.tokenizer_only = tokenizer_only |
| | |
| | self.template_parser = template_parser(meta_template) |
| | self.eos_token_id = None |
| | if meta_template and 'eos_token_id' in meta_template: |
| | self.eos_token_id = meta_template['eos_token_id'] |
| |
|
| | if isinstance(stop_words, str): |
| | stop_words = [stop_words] |
| | self.gen_params = dict( |
| | max_new_tokens=max_new_tokens, |
| | top_p=top_p, |
| | top_k=top_k, |
| | temperature=temperature, |
| | repetition_penalty=repetition_penalty, |
| | stop_words=stop_words) |
| |
|
| | def generate(self, inputs: Union[str, List[str]], **gen_params) -> str: |
| | """Generate results given a str (or list of) inputs. |
| | |
| | Args: |
| | inputs (Union[str, List[str]]): |
| | gen_params (dict): The input params for generation. |
| | |
| | Returns: |
| | Union[str, List[str]]: A (list of) generated strings. |
| | |
| | eg. |
| | batched = True |
| | if isinstance(inputs, str): |
| | inputs = [inputs] |
| | batched = False |
| | response = [''] |
| | if batched: |
| | return response |
| | return response[0] |
| | """ |
| | raise NotImplementedError |
| |
|
| | def stream_generate(self, inputs: str, **gen_params) -> List[str]: |
| | """Generate results as streaming given a str inputs. |
| | |
| | Args: |
| | inputs (str): |
| | gen_params (dict): The input params for generation. |
| | |
| | Returns: |
| | str: A generated string. |
| | """ |
| | raise NotImplementedError |
| |
|
| | def chat(self, |
| | inputs: Union[List[dict], List[List[dict]]], |
| | session_ids: Union[int, List[int]] = None, |
| | **gen_params): |
| | """Generate completion from a list of templates. |
| | |
| | Args: |
| | inputs (Union[List[dict], List[List[dict]]]): |
| | gen_params (dict): The input params for generation. |
| | Returns: |
| | """ |
| | if isinstance(inputs[0], list): |
| | _inputs = list() |
| | for msg in inputs: |
| | _inputs.append(self.template_parser(msg)) |
| | else: |
| | _inputs = self.template_parser(inputs) |
| | return self.generate(_inputs, **gen_params) |
| |
|
| | def stream_chat(self, inputs: List[dict], **gen_params): |
| | """Generate results as streaming given a list of templates. |
| | |
| | Args: |
| | inputs (Union[List[dict]): |
| | gen_params (dict): The input params for generation. |
| | Returns: |
| | """ |
| | raise NotImplementedError |
| |
|
| | def tokenize(self, prompts: Union[str, List[str], List[dict], |
| | List[List[dict]]]): |
| | """Tokenize the input prompts. |
| | |
| | Args: |
| | prompts(str | List[str]): user's prompt, or a batch prompts |
| | |
| | Returns: |
| | Tuple(numpy.ndarray, numpy.ndarray, numpy.ndarray): prompt's token |
| | ids, ids' length and requested output length |
| | """ |
| | raise NotImplementedError |
| |
|
| | def update_gen_params(self, **kwargs): |
| | gen_params = copy(self.gen_params) |
| | gen_params.update(kwargs) |
| | return gen_params |
| |
|
| |
|
| | class AsyncLLMMixin: |
| |
|
| | async def generate(self, |
| | inputs: Union[str, List[str]], |
| | session_ids: Union[int, List[int]] = None, |
| | **gen_params) -> str: |
| | """Generate results given a str (or list of) inputs. |
| | |
| | Args: |
| | inputs (Union[str, List[str]]): |
| | gen_params (dict): The input params for generation. |
| | |
| | Returns: |
| | Union[str, List[str]]: A (list of) generated strings. |
| | |
| | eg. |
| | batched = True |
| | if isinstance(inputs, str): |
| | inputs = [inputs] |
| | batched = False |
| | response = [''] |
| | if batched: |
| | return response |
| | return response[0] |
| | """ |
| | raise NotImplementedError |
| |
|
| | async def stream_generate(self, inputs: str, **gen_params) -> List[str]: |
| | """Generate results as streaming given a str inputs. |
| | |
| | Args: |
| | inputs (str): |
| | gen_params (dict): The input params for generation. |
| | |
| | Returns: |
| | str: A generated string. |
| | """ |
| | raise NotImplementedError |
| |
|
| | async def chat(self, |
| | inputs: Union[List[dict], List[List[dict]]], |
| | session_ids: Union[int, List[int]] = None, |
| | **gen_params): |
| | """Generate completion from a list of templates. |
| | |
| | Args: |
| | inputs (Union[List[dict], List[List[dict]]]): |
| | gen_params (dict): The input params for generation. |
| | Returns: |
| | """ |
| | if isinstance(inputs[0], list): |
| | _inputs = list() |
| | for msg in inputs: |
| | _inputs.append(self.template_parser(msg)) |
| | else: |
| | _inputs = self.template_parser(inputs) |
| | return await self.generate(_inputs, session_ids, **gen_params) |
| |
|
| | async def stream_chat(self, inputs: List[dict], **gen_params): |
| | """Generate results as streaming given a list of templates. |
| | |
| | Args: |
| | inputs (Union[List[dict]): |
| | gen_params (dict): The input params for generation. |
| | Returns: |
| | """ |
| | raise NotImplementedError |
| |
|
| | async def tokenize(self, prompts: Union[str, List[str], List[dict], |
| | List[List[dict]]]): |
| | """Tokenize the input prompts. |
| | |
| | Args: |
| | prompts(str | List[str]): user's prompt, or a batch prompts |
| | |
| | Returns: |
| | Tuple(numpy.ndarray, numpy.ndarray, numpy.ndarray): prompt's token |
| | ids, ids' length and requested output length |
| | """ |
| | raise NotImplementedError |
| |
|
| |
|
| | class AsyncBaseLLM(AsyncLLMMixin, BaseLLM): |
| | pass |
| |
|
