| import re | |
| from hat_splitter import HATSplitter as RustHATSplitter | |
| class HATSplitter: | |
| def __init__(self, special_token_dict: dict | None = None, max_word_size: int = 128): | |
| self.hat_splitter = RustHATSplitter() | |
| self.max_word_size = max_word_size | |
| self.special_token_dict = special_token_dict | |
| self.special_token_replace: dict[int, list[int]] = { | |
| token: list(text.encode("utf-8")) for text, token in self.special_token_dict.items() | |
| } | |
| self.special_token_pattern = ( | |
| re.compile(rf"({'|'.join(map(re.escape, special_token_dict.keys()))})") | |
| if special_token_dict | |
| else re.compile(r"(?!)") | |
| ) | |
| def encode(self, text: str) -> list[list[int]]: | |
| chunks = [] | |
| for str_chunk in self.special_token_pattern.split(text): | |
| if str_chunk: | |
| if str_chunk in self.special_token_dict: | |
| chunks.append([self.special_token_dict[str_chunk]]) | |
| else: | |
| chunks.extend(list(chunk) for chunk in self.hat_splitter.split_with_limit(str_chunk, self.max_word_size)) | |
| return chunks | |
| def decode(self, token_ids: list[int], errors: str = "replace", skip_special_tokens: bool = False) -> str: | |
| assert isinstance(token_ids, list), "token_ids must be a list" | |
| assert all(isinstance(token_id, int) for token_id in token_ids), "token_ids must be a list of integers" | |
| new_token_ids: list[int] | |
| if skip_special_tokens: | |
| new_token_ids = [token_id for token_id in token_ids if token_id not in self.special_token_replace] | |
| else: | |
| new_token_ids = [] | |
| for token in token_ids: | |
| if token in self.special_token_replace: | |
| new_token_ids.extend(self.special_token_replace[token]) | |
| else: | |
| new_token_ids.append(token) | |
| return bytes(new_token_ids).decode("utf-8", errors=errors) | |