|
|
|
import json
|
|
from typing import List, Optional, Tuple
|
|
|
|
from tokenizers import pre_tokenizers, processors
|
|
|
|
from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
|
|
from transformers.utils import logging
|
|
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}
|
|
|
|
|
|
class HelpingAITokenizerFast(PreTrainedTokenizerFast):
|
|
vocab_files_names = VOCAB_FILES_NAMES
|
|
model_input_names = ["input_ids", "attention_mask"]
|
|
|
|
def __init__(
|
|
self,
|
|
vocab_file=None,
|
|
merges_file=None,
|
|
tokenizer_file=None,
|
|
unk_token="<|endoftext|>",
|
|
bos_token="<|im_start|>",
|
|
eos_token="<|im_end|>",
|
|
pad_token="<|im_end|>",
|
|
add_bos_token=False,
|
|
add_eos_token=False,
|
|
add_prefix_space=False,
|
|
**kwargs,
|
|
):
|
|
super().__init__(
|
|
vocab_file,
|
|
merges_file,
|
|
tokenizer_file=tokenizer_file,
|
|
unk_token=unk_token,
|
|
bos_token=bos_token,
|
|
eos_token=eos_token,
|
|
pad_token=pad_token,
|
|
add_bos_token=add_bos_token,
|
|
add_eos_token=add_eos_token,
|
|
add_prefix_space=add_prefix_space,
|
|
**kwargs,
|
|
)
|
|
|
|
self._add_bos_token = add_bos_token
|
|
self._add_eos_token = add_eos_token
|
|
self.update_post_processor()
|
|
|
|
pre_tok_state = json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__())
|
|
if pre_tok_state.get("add_prefix_space", add_prefix_space) != add_prefix_space:
|
|
pre_tok_class = getattr(pre_tokenizers, pre_tok_state.pop("type"))
|
|
pre_tok_state["add_prefix_space"] = add_prefix_space
|
|
self.backend_tokenizer.pre_tokenizer = pre_tok_class(**pre_tok_state)
|
|
|
|
self.add_prefix_space = add_prefix_space
|
|
|
|
@property
|
|
def add_eos_token(self):
|
|
return self._add_eos_token
|
|
|
|
@property
|
|
def add_bos_token(self):
|
|
return self._add_bos_token
|
|
|
|
@add_eos_token.setter
|
|
def add_eos_token(self, value):
|
|
self._add_eos_token = value
|
|
self.update_post_processor()
|
|
|
|
@add_bos_token.setter
|
|
def add_bos_token(self, value):
|
|
self._add_bos_token = value
|
|
self.update_post_processor()
|
|
|
|
def update_post_processor(self):
|
|
bos = self.bos_token
|
|
bos_token_id = self.bos_token_id
|
|
if bos is None and self.add_bos_token:
|
|
raise ValueError("add_bos_token = True but bos_token = None")
|
|
|
|
eos = self.eos_token
|
|
eos_token_id = self.eos_token_id
|
|
if eos is None and self.add_eos_token:
|
|
raise ValueError("add_eos_token = True but eos_token = None")
|
|
|
|
single = f"{(bos+':0 ') if self.add_bos_token else ''}$A:0{(' '+eos+':0') if self.add_eos_token else ''}"
|
|
pair = f"{single}{(' '+bos+':1') if self.add_bos_token else ''} $B:1{(' '+eos+':1') if self.add_eos_token else ''}"
|
|
|
|
special_tokens = []
|
|
if self.add_bos_token:
|
|
special_tokens.append((bos, bos_token_id))
|
|
if self.add_eos_token:
|
|
special_tokens.append((eos, eos_token_id))
|
|
self._tokenizer.post_processor = processors.TemplateProcessing(
|
|
single=single, pair=pair, special_tokens=special_tokens
|
|
)
|
|
|
|
def get_special_tokens_mask(
|
|
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
|
|
) -> List[int]:
|
|
if already_has_special_tokens:
|
|
return super().get_special_tokens_mask(
|
|
token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
|
|
)
|
|
|
|
bos_token_id = [1] if self.add_bos_token else []
|
|
eos_token_id = [1] if self.add_eos_token else []
|
|
|
|
if token_ids_1 is None:
|
|
return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
|
|
return (
|
|
bos_token_id
|
|
+ ([0] * len(token_ids_0))
|
|
+ eos_token_id
|
|
+ bos_token_id
|
|
+ ([0] * len(token_ids_1))
|
|
+ eos_token_id
|
|
)
|
|
|
|
def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
|
|
bos_token_id = [self.bos_token_id] if self.add_bos_token else []
|
|
eos_token_id = [self.eos_token_id] if self.add_eos_token else []
|
|
|
|
output = bos_token_id + token_ids_0 + eos_token_id
|
|
|
|
if token_ids_1 is not None:
|
|
output = output + bos_token_id + token_ids_1 + eos_token_id
|
|
|
|
return output
|
|
|
|
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
|
|
files = self._tokenizer.model.save(save_directory, name=filename_prefix)
|
|
return tuple(files)
|
|
|
|
@property
|
|
def default_chat_template(self):
|
|
return "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
|
|
|
|
"""
|
|
Tool call template not added
|
|
""" |