HauserGroup/ApeTokenizer-SMILES

1"""Hugging Face-compatible tokenizer for APE molecular vocabularies.2 3This file is intentionally self-contained so it can be copied into a model repo4and loaded by ``AutoTokenizer.from_pretrained(..., trust_remote_code=True)``.5"""6 7import json8import os9import re10from collections.abc import Mapping11from collections import defaultdict12from pathlib import Path13from typing import Any, Literal14 15from transformers import PreTrainedTokenizer16 17 18Representation = Literal["SELFIES", "SMILES"]19 20VOCAB_FILES_NAMES = {21    "vocab_file": "vocab.json",22    "selfies_vocab_file": "selfies_vocab.json",23    "smiles_vocab_file": "smiles_vocab.json",24}25SELFIES_RE = re.compile(r"\[[^\]]+\]")26# Only the organic subset (B C N O P S F Cl Br I) may appear unbracketed in27# canonical SMILES; two-letter metals (Si, Se, Na, Mg, Al, Ca, Fe, Zn, ...) are28# always bracketed and matched by the leading \[[^\]]+\] branch. The previous29# pattern listed those metals as optional-second-letter alternatives (Si?, Na?,30# ...), which could match bare invalid single letters (L, M, A, Z) and was dead31# weight for valid input. Keep only Br?/Cl? (B, C, Br, Cl all valid bare).32SMILES_RE = re.compile(33    r"(\[[^\]]+\]|Br?|Cl?|"34    r"N|O|S|P|F|I|K|B|C|H|"35    r"b|c|n|o|s|p|"36    r"\%\d{2}|\d|"37    r"\(|\)|\.|=|#|-|\+|\\|/|:|~|@|\?|\*|\$)"38)39 40 41def _base_piece_count(token: str, representation: str) -> int:42    """Count primitive molecular pieces in a vocab token."""43    pieces = pre_tokenize_molecule(token, representation)44    return max(1, len(pieces))45 46 47def _max_vocab_piece_span(vocab: dict[str, int], representation: str) -> int:48    """Maximum number of primitive pieces covered by any non-special vocab token."""49    max_span = 150    for token in vocab:51        if token.startswith("<") and token.endswith(">"):52            continue53        max_span = max(max_span, _base_piece_count(token, representation))54    return max_span55 56 57def _coerce_vocab(vocab: Mapping[str, Any]) -> dict[str, int]:58    if not isinstance(vocab, Mapping):59        raise ValueError("Vocabulary must be a JSON object mapping token strings to integer IDs.")60    out = {str(token): int(idx) for token, idx in vocab.items()}61    if len(set(out.values())) != len(out):62        raise ValueError("Vocabulary token IDs must be unique.")63    return out64 65 66def _token_text(token: Any) -> str:67    return str(getattr(token, "content", token))68 69 70def _normalize_representation(representation: str) -> Representation:71    normalized = representation.upper()72    if normalized not in {"SELFIES", "SMILES"}:73        raise ValueError(f"representation must be 'SELFIES' or 'SMILES', got {representation!r}")74    return normalized  # type: ignore[return-value]75 76 77def _select_vocab_file(78    *,79    representation: Representation,80    vocab_file: str | os.PathLike[str] | None,81    selfies_vocab_file: str | os.PathLike[str] | None,82    smiles_vocab_file: str | os.PathLike[str] | None,83) -> str | os.PathLike[str] | None:84    if representation == "SELFIES" and selfies_vocab_file is not None:85        return selfies_vocab_file86    if representation == "SMILES" and smiles_vocab_file is not None:87        return smiles_vocab_file88    return vocab_file89 90 91def _pre_tokenize_selfies(molecule: str, *, strict: bool = True) -> list[str]:92    pieces = SELFIES_RE.findall(molecule)93 94    if strict and "".join(pieces) != molecule:95        raise ValueError(96            "Malformed SELFIES string contains unmatched text outside "97            f"bracketed SELFIES tokens: {molecule!r}"98        )99 100    return pieces101 102 103def pre_tokenize_molecule(104    molecule: str,105    representation: str,106    *,107    strict_selfies: bool = True,108) -> list[str]:109    active_representation = _normalize_representation(representation)110 111    if active_representation == "SELFIES":112        return _pre_tokenize_selfies(molecule, strict=strict_selfies)113 114    tokens: list[str] = []115    cursor = 0116 117    for match in SMILES_RE.finditer(molecule):118        if match.start() > cursor:119            tokens.extend(molecule[cursor : match.start()])120 121        tokens.append(match.group(0))122        cursor = match.end()123 124    if cursor < len(molecule):125        tokens.extend(molecule[cursor:])126 127    return [token for token in tokens if token and not token.isspace()]128 129 130def ape_tokenize(131    text: str,132    vocab: dict[str, int],133    representation: str,134    unk_token: str = "<unk>",135    max_piece_span: int | None = None,136) -> list[str]:137    """Segment a molecule against the APE vocabulary by greedy longest match.138 139    Note this is *not* a replay of the training merges in learned order: train()140    learns which substrings become vocab entries, but decoding here just takes141    the longest vocab token at each position (up to ``max_piece_span`` pieces).142    The two can disagree on segmentation. That is fine and intended — both143    pretraining and fine-tuning encode through this same function, so the model144    only ever sees greedy-longest-match output and stays internally consistent.145    The learned merge *order* is intentionally discarded; only the vocab set is146    used at inference.147    """148    # A single malformed SELFIES (stray text outside bracket tokens) must not149    # crash encoding. Map the whole string to <unk> so it stays detectable via150    # the validator's unk_rate gate instead of raising mid-batch.151    try:152        pieces = pre_tokenize_molecule(text, representation)153    except ValueError:154        return [unk_token]155    if not pieces:156        return [unk_token]157 158    if max_piece_span is None:159        max_piece_span = _max_vocab_piece_span(vocab, representation)160 161    n = len(pieces)162    tokens: list[str] = []163    append_token = tokens.append164    vocab_contains = vocab.__contains__165    join_pieces = "".join166    i = 0167 168    while i < n:169        upper = min(n, i + max_piece_span)170 171        for j in range(upper, i, -1):172            candidate = join_pieces(pieces[i:j])173            if vocab_contains(candidate):174                append_token(candidate)175                i = j176                break177        else:178            append_token(unk_token)179            i += 1180 181    return tokens182 183 184class APEPreTrainedTokenizer(PreTrainedTokenizer):185    """Hugging Face tokenizer backend for APE molecular tokenization. (Not fast)"""186 187    vocab_files_names = VOCAB_FILES_NAMES188    model_input_names = ["input_ids", "attention_mask"]189 190    def __init__(191        self,192        vocab_file: str | os.PathLike[str] | None = None,193        selfies_vocab_file: str | os.PathLike[str] | None = None,194        smiles_vocab_file: str | os.PathLike[str] | None = None,195        vocab: dict[str, Any] | None = None,196        representation: str = "SELFIES",197        bos_token: str = "<s>",198        eos_token: str = "</s>",199        unk_token: str = "<unk>",200        pad_token: str = "<pad>",201        mask_token: str = "<mask>",202        model_max_length: int = 256,203        **kwargs,204    ) -> None:205        self.representation = _normalize_representation(representation)206        active_vocab_file = _select_vocab_file(207            representation=self.representation,208            vocab_file=vocab_file,209            selfies_vocab_file=selfies_vocab_file,210            smiles_vocab_file=smiles_vocab_file,211        )212 213        if vocab is None:214            if active_vocab_file is None:215                vocab = {216                    bos_token: 0,217                    pad_token: 1,218                    eos_token: 2,219                    unk_token: 3,220                    mask_token: 4,221                }222            else:223                with open(active_vocab_file, encoding="utf-8") as f:224                    vocab = json.load(f)225 226        if vocab is None:227            raise ValueError("Loaded vocabulary is None.")228 229        self.vocab_file = str(active_vocab_file) if active_vocab_file is not None else None230        self.selfies_vocab_file = (231            str(selfies_vocab_file) if selfies_vocab_file is not None else None232        )233        self.smiles_vocab_file = str(smiles_vocab_file) if smiles_vocab_file is not None else None234        self.vocab = _coerce_vocab(vocab)235        self._require_special_tokens(236            bos_token=bos_token,237            eos_token=eos_token,238            unk_token=unk_token,239            pad_token=pad_token,240            mask_token=mask_token,241        )242        self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}243        self.vocabulary_frequency: dict[str, int] = {}244        self.pair_counts: dict[tuple[str, str], int] = {}245        self._max_piece_span = _max_vocab_piece_span(self.vocab, self.representation)246 247        super().__init__(248            bos_token=bos_token,249            eos_token=eos_token,250            unk_token=unk_token,251            pad_token=pad_token,252            mask_token=mask_token,253            model_max_length=model_max_length,254            representation=self.representation,255            **kwargs,256        )257 258    @property259    def vocab_size(self) -> int:260        return len(self.vocab)261 262    @property263    def vocabulary(self) -> dict[str, int]:264        """Legacy alias for callers that previously used APETokenizer."""265        return self.vocab266 267    @vocabulary.setter268    def vocabulary(self, value: dict[str, int]) -> None:269        self.vocab = _coerce_vocab(value)270        self.update_reverse_vocabulary()271        self._refresh_tokenization_cache()272 273    @property274    def special_tokens(self) -> dict[str, int]:275        bos_token = str(self.bos_token)276        pad_token = str(self.pad_token)277        eos_token = str(self.eos_token)278        unk_token = str(self.unk_token)279        mask_token = str(self.mask_token)280        return {281            bos_token: self._convert_token_to_id(bos_token),282            pad_token: self._convert_token_to_id(pad_token),283            eos_token: self._convert_token_to_id(eos_token),284            unk_token: self._convert_token_to_id(unk_token),285            mask_token: self._convert_token_to_id(mask_token),286        }287 288    @special_tokens.setter289    def special_tokens(self, value: dict[str, int]) -> None:290        for token, token_id in value.items():291            self.vocab.setdefault(str(token), int(token_id))292        self.vocab = _coerce_vocab(self.vocab)293        self.update_reverse_vocabulary()294        self._refresh_tokenization_cache()295 296    def get_vocab(self) -> dict[str, int]:297        return dict(self.vocab)298 299    def update_reverse_vocabulary(self) -> None:300        self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}301 302    def _refresh_tokenization_cache(self) -> None:303        self._max_piece_span = _max_vocab_piece_span(self.vocab, self.representation)304 305    def _require_special_tokens(306        self,307        *,308        bos_token: str,309        eos_token: str,310        unk_token: str,311        pad_token: str,312        mask_token: str,313    ) -> None:314        missing = [315            token_text316            for token in [bos_token, eos_token, unk_token, pad_token, mask_token]317            if (token_text := _token_text(token)) not in self.vocab318        ]319        if missing:320            raise ValueError(f"Vocabulary is missing required special tokens: {missing}")321 322    def pre_tokenize(self, molecule: str, representation: str | None = None) -> list[str]:323        return pre_tokenize_molecule(molecule, representation or self.representation)324 325    def _tokenize(self, text: str, **kwargs) -> list[str]:326 327        return ape_tokenize(328            text,329            vocab=self.vocab,330            representation=self.representation,331            unk_token=str(self.unk_token),332            max_piece_span=self._max_piece_span,333        )334 335    def encode_molecule(336        self,337        text: str,338        add_special_tokens: bool = True,339        max_length: int | None = None,340        truncation: bool = True,341    ) -> list[int]:342        """Fast molecular encode path avoiding generic Hugging Face tokenizer overhead."""343 344        tokens = self._tokenize(text)345 346        ids = [self._convert_token_to_id(token) for token in tokens]347 348        if add_special_tokens:349            ids = self.build_inputs_with_special_tokens(ids)350 351        if max_length is not None and truncation:352            ids = ids[:max_length]353 354        return ids355 356    def _convert_token_to_id(self, token: str) -> int:357        return self.vocab.get(token, self.vocab[str(self.unk_token)])358 359    def _convert_id_to_token(self, index: int) -> str:360        return self.ids_to_tokens.get(int(index), str(self.unk_token))361 362    def convert_tokens_to_string(self, tokens: list[str]) -> str:363        return "".join(tokens)364 365    def _required_special_token_id(366        self,367        token_value: int | list[int] | str | list[str] | None,368        token_name: str,369    ) -> int:370        if token_value is None:371            raise ValueError(f"{token_name} must be set.")372        if isinstance(token_value, int):373            return token_value374        if isinstance(token_value, str):375            return self._convert_token_to_id(token_value)376        if len(token_value) == 1:377            only_value = token_value[0]378            if isinstance(only_value, int):379                return only_value380            if isinstance(only_value, str):381                return self._convert_token_to_id(only_value)382        raise ValueError(f"{token_name} must resolve to a single token id.")383 384    def build_inputs_with_special_tokens(385        self,386        token_ids_0: list[int],387        token_ids_1: list[int] | None = None,388    ) -> list[int]:389        bos_id = self._required_special_token_id(self.bos_token, "bos_token")390        eos_id = self._required_special_token_id(self.eos_token, "eos_token")391        if token_ids_1 is None:392            return [bos_id, *token_ids_0, eos_id]393        return [bos_id, *token_ids_0, eos_id, *token_ids_1, eos_id]394 395    def create_token_type_ids_from_sequences(396        self,397        token_ids_0: list[int],398        token_ids_1: list[int] | None = None,399    ) -> list[int]:400        return [0] * len(self.build_inputs_with_special_tokens(token_ids_0, token_ids_1))401 402    def pad(403        self,404        encoded_inputs: Any,405        padding: Any = True,406        max_length: int | None = None,407        pad_to_multiple_of: int | None = None,408        padding_side: str | None = None,409        return_attention_mask: bool | None = None,410        return_tensors: Any = None,411        verbose: bool = True,412    ):413        padding_enabled = padding not in (False, "do_not_pad")414        if (415            padding_enabled416            and isinstance(encoded_inputs, list)417            and any("labels" in item for item in encoded_inputs)418        ):419            target_length = max(420                len(item.get("input_ids", item.get("labels", []))) for item in encoded_inputs421            )422            if padding == "max_length" and max_length is not None:423                target_length = max_length424 425            if pad_to_multiple_of and target_length % pad_to_multiple_of:426                target_length = ((target_length // pad_to_multiple_of) + 1) * pad_to_multiple_of427 428            padded_inputs = []429            for item in encoded_inputs:430                item = dict(item)431                labels = list(item.get("labels", []))432                pad_len = max(0, target_length - len(labels))433                if pad_len:434                    label_padding = [-100] * pad_len435                    if self.padding_side == "left":436                        labels = label_padding + labels437                    else:438                        labels = labels + label_padding439                    item["labels"] = labels440                padded_inputs.append(item)441            encoded_inputs = padded_inputs442 443        return super().pad(444            encoded_inputs,445            padding=padding,446            max_length=max_length,447            pad_to_multiple_of=pad_to_multiple_of,448            padding_side=padding_side,449            return_attention_mask=return_attention_mask,450            return_tensors=return_tensors,451            verbose=verbose,452        )453 454    def save_vocabulary(455        self,456        save_directory: str,457        filename_prefix: str | None = None,458    ) -> tuple[str, ...]:459        if not os.path.isdir(save_directory):460            raise ValueError(f"Vocabulary path ({save_directory}) should be a directory.")461 462        vocab_file = Path(save_directory) / (463            f"{filename_prefix}-vocab.json" if filename_prefix else "vocab.json"464        )465        with vocab_file.open("w", encoding="utf-8") as f:466            json.dump(self.vocab, f, ensure_ascii=False, indent=4)467        return (str(vocab_file),)468 469    def add_tokens_to_vocabulary(self, tokens: list[str]) -> int:470        """Add tokens to the tokenizer vocabulary if they are not already present.471 472        This is intended for forcing coverage of rare valid molecular primitive473        symbols, especially SELFIES bracket tokens, after APE merge training.474        """475 476        if not tokens:477            return 0478 479        next_id = max(self.vocab.values(), default=-1) + 1480        added = 0481 482        for token in tokens:483            token = str(token).strip()484            if not token:485                continue486            if token in self.vocab:487                continue488 489            self.vocab[token] = next_id490            next_id += 1491            added += 1492 493        if added:494            self.update_reverse_vocabulary()495            self._refresh_tokenization_cache()496 497        return added498 499    def save_pretrained(self, save_directory: str | os.PathLike[str], *args, **kwargs):500        saved_files = super().save_pretrained(save_directory, *args, **kwargs)501        save_path = Path(save_directory)502 503        special_tokens_map = {504            "bos_token": str(self.bos_token),505            "eos_token": str(self.eos_token),506            "unk_token": str(self.unk_token),507            "pad_token": str(self.pad_token),508            "mask_token": str(self.mask_token),509        }510        with (save_path / "special_tokens_map.json").open("w", encoding="utf-8") as f:511            json.dump(special_tokens_map, f, ensure_ascii=False, indent=2)512 513        tokenizer_config_path = save_path / "tokenizer_config.json"514        if tokenizer_config_path.exists():515            with tokenizer_config_path.open(encoding="utf-8") as f:516                tokenizer_config = json.load(f)517        else:518            tokenizer_config = {}519        tokenizer_config.pop("tokenizer_class", None)520        tokenizer_config.update(521            {522                "representation": self.representation,523                "model_max_length": self.model_max_length,524                "auto_map": {525                    "AutoTokenizer": [526                        "tokenization_ape.APEPreTrainedTokenizer",527                        None,528                    ],529                },530            }531        )532        with tokenizer_config_path.open("w", encoding="utf-8") as f:533            json.dump(tokenizer_config, f, ensure_ascii=False, indent=2)534 535        return saved_files536 537    def save_vocabulary_file(self, file_path: str | os.PathLike[str]) -> None:538        path = Path(file_path)539        path.parent.mkdir(parents=True, exist_ok=True)540        freq_path = path.with_name(f"{path.stem}_freq.json")541 542        with path.open("w", encoding="utf-8") as f:543            json.dump(self.vocab, f, ensure_ascii=False, indent=4)544        with freq_path.open("w", encoding="utf-8") as f:545            json.dump(self.vocabulary_frequency, f, ensure_ascii=False, indent=4)546 547    def load_vocabulary_file(548        self,549        file_path: str | os.PathLike[str],550        representation: str | None = None,551    ) -> None:552        if representation is not None:553            self.representation = _normalize_representation(representation)554        with open(file_path, encoding="utf-8") as f:555            vocab = json.load(f)556        self.vocab = _coerce_vocab(vocab)557        self._require_special_tokens(558            bos_token=str(self.bos_token),559            eos_token=str(self.eos_token),560            unk_token=str(self.unk_token),561            pad_token=str(self.pad_token),562            mask_token=str(self.mask_token),563        )564        self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}565        self._refresh_tokenization_cache()566 567    def train(568        self,569        corpus,570        type: str = "selfies",571        representation: str | None = None,572        max_vocab_size: int = 5000,573        min_freq_for_merge: int = 2000,574        max_merge_pieces: int | None = 8,575        save_checkpoint: bool = False,576        checkpoint_path: str = "checkpoint",577        checkpoint_interval: int = 500,578    ) -> None:579        import warnings580 581        new_rep = _normalize_representation(representation or type)582        if new_rep != self.representation:583            warnings.warn(584                f"train() representation={new_rep!r} differs from tokenizer "585                f"representation={self.representation!r}. Overwriting.",586                UserWarning,587                stacklevel=2,588            )589        self.representation = new_rep590 591        if not corpus:592            raise ValueError("Cannot train APE tokenizer on an empty corpus.")593 594        print(f"Pretokenizing {self.representation}...", flush=True)595        tokenized_corpus = []596        vocabulary_frequency: defaultdict[str, int] = defaultdict(int)597        saw_tokens = False598        skipped_malformed = 0599 600        for sentence in corpus:601            # One malformed row must not abort a multi-hour training run. Skip and602            # count it; surface the total so a corrupt corpus is still visible.603            try:604                tokens = self.pre_tokenize(str(sentence))605            except ValueError:606                skipped_malformed += 1607                continue608            if not tokens:609                continue610            saw_tokens = True611            for token in tokens:612                vocabulary_frequency[token] += 1613            if len(tokens) > 1:614                tokenized_corpus.append(tokens)615        if skipped_malformed:616            print(f"Skipped {skipped_malformed} malformed sequences", flush=True)617        print(618            f"Pretokenization complete, found {len(vocabulary_frequency)} tokens",619            flush=True,620        )621 622        if not saw_tokens:623            raise ValueError("Cannot train APE tokenizer on an empty corpus.")624 625        pre_tokens_counts = len(vocabulary_frequency)626        merged_counter = len(vocabulary_frequency) + 1627        if save_checkpoint and checkpoint_interval <= 0:628            raise ValueError(629                "checkpoint_interval must be positive when save_checkpoint is enabled."630            )631        checkpoint_increment = checkpoint_interval632        batch = checkpoint_interval + pre_tokens_counts633        piece_count_cache: dict[str, int] = {}634 635        def merged_piece_count(token: str) -> int:636            count = piece_count_cache.get(token)637            if count is None:638                count = _base_piece_count(token, self.representation)639                piece_count_cache[token] = count640            return count641 642        def get_most_common_pair(tokenized):643            pair_counts: defaultdict[tuple[str, str], int] = defaultdict(int)644            for tokens in tokenized:645                for i in range(len(tokens) - 1):646                    pair = (tokens[i], tokens[i + 1])647 648                    if max_merge_pieces is not None:649                        merged_candidate = pair[0] + pair[1]650                        if merged_piece_count(merged_candidate) > max_merge_pieces:651                            continue652 653                    pair_counts[pair] += 1654 655            if not pair_counts:656                return ("", ""), 0657 658            most_common_pair = ("", "")659            most_common_frequency = 0660            for pair, count in pair_counts.items():661                if count > most_common_frequency:662                    most_common_pair = pair663                    most_common_frequency = count664            return most_common_pair, most_common_frequency665 666        while True:667            if save_checkpoint and len(vocabulary_frequency) >= batch:668                self.vocabulary_frequency = dict(vocabulary_frequency)669                self.vocab = {670                    **{671                        str(self.bos_token): 0,672                        str(self.pad_token): 1,673                        str(self.eos_token): 2,674                        str(self.unk_token): 3,675                        str(self.mask_token): 4,676                    },677                    **{678                        word: idx679                        for idx, word in enumerate(680                            vocabulary_frequency.keys(),681                            start=5,682                        )683                    },684                }685                self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}686                self._refresh_tokenization_cache()687                checkpoint_dir = Path(checkpoint_path)688                checkpoint_dir.mkdir(parents=True, exist_ok=True)689                self.save_vocabulary_file(checkpoint_dir / f"checkpoint_{batch}.json")690                self.save_pretrained(str(checkpoint_dir / f"checkpoint_{batch}"))691                print(f"Checkpoint saved at {checkpoint_dir}/checkpoint_{batch}.json")692                batch += checkpoint_increment693 694            if len(vocabulary_frequency) >= max_vocab_size:695                print("Max vocabulary achieved", flush=True)696                break697 698            if not tokenized_corpus:699                print("No more mergeable pairs", flush=True)700                break701 702            most_common_pair, freq = get_most_common_pair(tokenized_corpus)703            if freq < min_freq_for_merge:704                print("Not enough frequency found", flush=True)705                break706 707            if not most_common_pair[0] or not most_common_pair[1]:708                print("No valid merge pair found", flush=True)709                break710 711            left_token, right_token = most_common_pair712            merged_word = left_token + right_token713            if merged_word not in vocabulary_frequency:714                print(715                    f"New merge found: {merged_word} {merged_counter}/{max_vocab_size} "716                    f"{round(merged_counter / max_vocab_size * 100, 2)}%",717                    flush=True,718                )719                merged_counter += 1720            # Each merged occurrence consumes one left + one right piece, so debit721            # both constituents to keep vocabulary_frequency (the *_freq.json722            # diagnostic) an accurate post-merge count. Keys are never removed —723            # a primitive merged to zero must stay in vocab for coverage.724            vocabulary_frequency[merged_word] += freq725            vocabulary_frequency[left_token] = max(0, vocabulary_frequency[left_token] - freq)726            vocabulary_frequency[right_token] = max(0, vocabulary_frequency[right_token] - freq)727 728            new_tokenized_corpus = []729            append_seq = new_tokenized_corpus.append730            for tokens in tokenized_corpus:731                token_count = len(tokens)732 733                # Fast path: a sequence with no adjacent (left, right) is734                # unchanged by this merge. Keep the existing list by reference735                # instead of reallocating + re-appending every token. Most736                # sequences are untouched per merge, so this avoids the bulk of737                # the per-iteration allocation without altering the output.738                has_pair = any(739                    tokens[i] == left_token and tokens[i + 1] == right_token740                    for i in range(token_count - 1)741                )742                if not has_pair:743                    append_seq(tokens)744                    continue745 746                new_tokens = []747                append_token = new_tokens.append748                i = 0749                while i < token_count:750                    if (751                        i < token_count - 1752                        and tokens[i] == left_token753                        and tokens[i + 1] == right_token754                    ):755                        append_token(merged_word)756                        i += 2757                    else:758                        append_token(tokens[i])759                        i += 1760 761                if len(new_tokens) > 1:762                    append_seq(new_tokens)763 764            tokenized_corpus = new_tokenized_corpus765 766        self.vocabulary_frequency = dict(vocabulary_frequency)767        self.vocab = {768            str(self.bos_token): 0,769            str(self.pad_token): 1,770            str(self.eos_token): 2,771            str(self.unk_token): 3,772            str(self.mask_token): 4,773            **{word: idx for idx, word in enumerate(vocabulary_frequency.keys(), start=5)},774        }775 776        self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}777        self._refresh_tokenization_cache()778 779    def train_from_iterator(self, iterator, *args, **kwargs) -> None:780        raise NotImplementedError("train_from_iterator is not implemented for APE")781 782 783APEPreTrainedTokenizer.register_for_auto_class("AutoTokenizer")784