|
|
|
|
|
|
|
|
from copy import deepcopy |
|
|
from dataclasses import dataclass, field |
|
|
from typing import List, Optional, Type, Union |
|
|
|
|
|
from transformers import PreTrainedTokenizerBase |
|
|
|
|
|
from .base import Template |
|
|
from .utils import Prompt, Word |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class TemplateMeta: |
|
|
""" |
|
|
Examples: |
|
|
chatml (with bos): |
|
|
prefix: <s> |
|
|
prompt: <|im_start|>user\n{{QUERY}}<|im_end|>\n<|im_start|>assistant\n |
|
|
chat_sep: <|im_end|>\n |
|
|
suffix: <|im_end|> |
|
|
system_prefix: <s><|im_start|>system\n{{SYSTEM}}<|im_end|>\n |
|
|
|
|
|
<s><|im_start|>system # prefix or system_prefix |
|
|
{{SYSTEM}}<|im_end|> |
|
|
<|im_start|>user # prompt |
|
|
{{QUERY}}<|im_end|> |
|
|
<|im_start|>assistant |
|
|
{{RESPONSE}}<|im_end|> # chat_sep |
|
|
<|im_start|>user # prompt |
|
|
{{QUERY}}<|im_end|> |
|
|
<|im_start|>assistant |
|
|
{{RESPONSE}}<|im_end|> # suffix |
|
|
""" |
|
|
template_type: str |
|
|
prefix: Prompt |
|
|
prompt: Prompt |
|
|
chat_sep: Optional[Prompt] |
|
|
suffix: Prompt = field(default_factory=lambda: [['eos_token_id']]) |
|
|
template_cls: Type[Template] = Template |
|
|
system_prefix: Optional[Prompt] = None |
|
|
default_system: Optional[str] = None |
|
|
response_prefix: str = '' |
|
|
|
|
|
auto_add_bos: bool = False |
|
|
stop_words: List[Word] = field(default_factory=list) |
|
|
agent_template: str = 'react_en' |
|
|
|
|
|
def to_generate_template_meta(self) -> 'TemplateMeta': |
|
|
self = deepcopy(self) |
|
|
return TemplateMeta( |
|
|
self.template_type, |
|
|
prefix=[], |
|
|
prompt=['{{QUERY}}'], |
|
|
chat_sep=None, |
|
|
template_cls=self.template_cls, |
|
|
auto_add_bos=True, |
|
|
stop_words=self.stop_words, |
|
|
) |
|
|
|
|
|
@staticmethod |
|
|
def _has_system(prefix_or_prompt: Prompt) -> bool: |
|
|
return any(['{{SYSTEM}}' in p for p in prefix_or_prompt]) |
|
|
|
|
|
@staticmethod |
|
|
def _replace_system(prefix: Prompt) -> Prompt: |
|
|
return [p.replace('{{SYSTEM}}', '') for p in prefix if isinstance(p, str)] |
|
|
|
|
|
def _check_template_meta(self): |
|
|
|
|
|
for x in [self.prefix, self.prompt, self.suffix]: |
|
|
assert isinstance(x, list) |
|
|
for x in [self.chat_sep, self.system_prefix]: |
|
|
assert x is None or isinstance(x, list) |
|
|
|
|
|
def __post_init__(self): |
|
|
|
|
|
if self._has_system(self.prefix): |
|
|
assert self.system_prefix is None, 'The prefix already contains {{SYSTEM}}.' |
|
|
self.system_prefix = self.prefix |
|
|
self.prefix = self._replace_system(self.prefix) |
|
|
|
|
|
self.is_post_system = self._has_system(self.prompt) |
|
|
if self.is_post_system: |
|
|
self.prompt = [context for context in self.prompt if '{{SYSTEM}}' not in context] |
|
|
self.system_prompt = self.prompt |
|
|
|
|
|
if self.system_prefix is None and not self.is_post_system: |
|
|
self.support_system = False |
|
|
else: |
|
|
self.support_system = True |
|
|
self.check_system(self.default_system) |
|
|
|
|
|
self.support_multi_round = self.chat_sep is not None |
|
|
|
|
|
@staticmethod |
|
|
def _token_attr_to_id(tokenizer: PreTrainedTokenizerBase, value: Optional[Prompt]) -> Optional[Prompt]: |
|
|
"""Turn `eos_token_id` to token id |
|
|
|
|
|
e.g. [['eos_token_id']] -> [[2]] |
|
|
""" |
|
|
if value is None: |
|
|
return None |
|
|
res_value = [] |
|
|
for v in value: |
|
|
if isinstance(v, list): |
|
|
v = [getattr(tokenizer, sub_v) if isinstance(sub_v, str) else sub_v for sub_v in v] |
|
|
res_value.append(v) |
|
|
return res_value |
|
|
|
|
|
def init(self, tokenizer: PreTrainedTokenizerBase) -> None: |
|
|
for key in ['prefix', 'prompt', 'chat_sep', 'suffix', 'system_prefix']: |
|
|
value = getattr(self, key) |
|
|
value = self._token_attr_to_id(tokenizer, value) |
|
|
setattr(self, key, value) |
|
|
|
|
|
if self.suffix and self.suffix[-1] not in self.stop_words: |
|
|
self.stop_words.append(self.suffix[-1]) |
|
|
if tokenizer.eos_token not in self.stop_words: |
|
|
self.stop_words.append(tokenizer.eos_token) |
|
|
|
|
|
self.stop_token_id = tokenizer.eos_token_id |
|
|
if self.suffix: |
|
|
suffix_tokens = self.suffix[-1] |
|
|
if isinstance(suffix_tokens, str): |
|
|
stop_token_id = tokenizer.convert_tokens_to_ids(suffix_tokens) |
|
|
elif isinstance(suffix_tokens, list) and len(suffix_tokens) == 1: |
|
|
stop_token_id = suffix_tokens[0] |
|
|
else: |
|
|
stop_token_id = None |
|
|
if stop_token_id is not None: |
|
|
self.stop_token_id = stop_token_id |
|
|
|
|
|
def check_system(self, system: Optional[str]) -> None: |
|
|
if system is not None: |
|
|
assert self.support_system, ( |
|
|
f'The template does not support `system`, template_type: {self.template_type}, system: {system}') |
|
|
|
