|
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 |
|
|
