tokenization_molparser_mobile.py

"""MolParser Mobile tokenizer for Hugging Face Hub remote loading."""

from __future__ import annotations

import json
import os
import re
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Sequence, Union

from huggingface_hub import hf_hub_download
from transformers import PreTrainedTokenizer


TOKENIZER_CONFIG_NAME = "tokenizer_config.json"
VOCAB_NAME = "vocab.txt"


class MolParserTokenizer(PreTrainedTokenizer):
    model_input_names = ["input_ids", "attention_mask"]
    padding_side = "right"

    def __init__(
        self,
        vocab_list: Optional[List[str]] = None,
        special_tokens: Optional[Dict[str, str]] = None,
        additional_special_tokens: Optional[Sequence[str]] = None,
        **kwargs,
    ):
        if vocab_list is None:
            vocab_list = []
        if special_tokens is None:
            special_tokens = {}
        if additional_special_tokens is None:
            additional_special_tokens = []

        self.special_tokens = {
            "cls_token": special_tokens.get("cls_token", "[CLS]"),
            "pad_token": special_tokens.get("pad_token", "[PAD]"),
            "sep_token": special_tokens.get("sep_token", "[SEP]"),
            "unk_token": special_tokens.get("unk_token", "[UNK]"),
        }
        self.additional_special_tokens = list(dict.fromkeys(additional_special_tokens))
        self.vocab_list = list(vocab_list)
        all_tokens = self._build_full_vocab(self.vocab_list)
        self.vocab = {token: idx for idx, token in enumerate(all_tokens)}
        self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}
        self._decode_skip_tokens = set(self.special_tokens.values())
        self._compile_pattern()

        super().__init__(
            cls_token=self.special_tokens["cls_token"],
            pad_token=self.special_tokens["pad_token"],
            sep_token=self.special_tokens["sep_token"],
            unk_token=self.special_tokens["unk_token"],
            bos_token=self.special_tokens["cls_token"],
            eos_token=self.special_tokens["sep_token"],
            additional_special_tokens=self.additional_special_tokens,
            **kwargs,
        )

    def _build_full_vocab(self, vocab_list: Sequence[str]) -> List[str]:
        ordered_tokens: List[str] = []
        for token in list(self.special_tokens.values()) + list(vocab_list) + list(self.additional_special_tokens):
            if token not in ordered_tokens:
                ordered_tokens.append(token)
        return ordered_tokens

    def _compile_pattern(self) -> None:
        multi_char_tokens = sorted(self.vocab.keys(), key=len, reverse=True)
        pattern = "(" + "|".join(re.escape(token) for token in multi_char_tokens) + "|.)"
        self.pattern = re.compile(pattern)

    @property
    def vocab_size(self) -> int:
        return len(self.vocab)

    def __len__(self) -> int:
        return len(self.vocab)

    def get_vocab(self) -> Dict[str, int]:
        return dict(self.vocab)

    def _tokenize(self, text: str) -> List[str]:
        return [token for token in self.pattern.findall(str(text)) if token]

    def tokenize(self, text: str, **kwargs) -> List[str]:
        return self._tokenize(text)

    def _convert_token_to_id(self, token: str) -> int:
        return self.vocab.get(token, self.unk_token_id)

    def _convert_id_to_token(self, index: int) -> str:
        return self.ids_to_tokens.get(int(index), self.unk_token)

    def convert_tokens_to_string(self, tokens: Sequence[str]) -> str:
        return "".join(tokens)

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        if token_ids_1 is None:
            return list(token_ids_0)
        return list(token_ids_0) + list(token_ids_1)

    def encode(self, text: str, add_special_tokens: bool = False, **kwargs) -> List[int]:
        token_ids = [self._convert_token_to_id(token) for token in self._tokenize(text)]
        if add_special_tokens:
            return [self.bos_token_id] + token_ids + [self.eos_token_id]
        return token_ids

    def decode(self, token_ids: Iterable[int], skip_special_tokens: bool = False, **kwargs) -> str:
        tokens = [self._convert_id_to_token(idx) for idx in token_ids]
        if skip_special_tokens:
            # Match deploy/tokenizer.py: keep MolParser business tokens such as
            # <sep>, <a>, </a>, <r>, </r>, <c>, </c>, and |Sg:n|.
            tokens = [token for token in tokens if token not in self._decode_skip_tokens]
        return "".join(tokens)

    def batch_encode(self, texts: Sequence[str], add_special_tokens: bool = False) -> List[List[int]]:
        return [self.encode(text, add_special_tokens=add_special_tokens) for text in texts]

    def batch_decode(
        self,
        sequences: Sequence[Sequence[int]],
        skip_special_tokens: bool = False,
        **kwargs,
    ) -> List[str]:
        return [self.decode(ids, skip_special_tokens=skip_special_tokens, **kwargs) for ids in sequences]

    def to_dict(self) -> Dict[str, object]:
        return {
            "vocab_list": self.vocab_list,
            "special_tokens": self.special_tokens,
            "additional_special_tokens": self.additional_special_tokens,
            "tokenizer_class": self.__class__.__name__,
            "auto_map": {
                "AutoTokenizer": [
                    "tokenization_molparser_mobile.MolParserTokenizer",
                    None,
                ]
            },
        }

    @classmethod
    def from_dict(cls, config: Dict[str, object]) -> "MolParserTokenizer":
        return cls(
            vocab_list=list(config["vocab_list"]),
            special_tokens=dict(config["special_tokens"]),
            additional_special_tokens=list(config.get("additional_special_tokens", [])),
        )

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None):
        path = Path(save_directory)
        path.mkdir(parents=True, exist_ok=True)
        name = f"{filename_prefix}-{VOCAB_NAME}" if filename_prefix else VOCAB_NAME
        vocab_path = path / name
        vocab_path.write_text("\n".join(self.vocab_list) + "\n", encoding="utf-8")
        return (str(vocab_path),)

    def save_pretrained(self, save_directory: str, **kwargs):
        os.makedirs(save_directory, exist_ok=True)
        config_path = os.path.join(save_directory, TOKENIZER_CONFIG_NAME)
        with open(config_path, "w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, ensure_ascii=False, indent=2)
        vocab_files = self.save_vocabulary(save_directory)
        return (config_path, *vocab_files)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs) -> "MolParserTokenizer":
        config_path = Path(pretrained_model_name_or_path)
        if config_path.is_dir():
            config_path = config_path / TOKENIZER_CONFIG_NAME
        elif config_path.is_file():
            pass
        else:
            config_path = Path(
                hf_hub_download(
                    repo_id=str(pretrained_model_name_or_path),
                    filename=TOKENIZER_CONFIG_NAME,
                    repo_type=kwargs.get("repo_type"),
                    revision=kwargs.get("revision"),
                    cache_dir=kwargs.get("cache_dir"),
                    token=kwargs.get("token"),
                    local_files_only=kwargs.get("local_files_only", False),
                )
            )
        with open(config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
        return cls.from_dict(config)


__all__ = ["MolParserTokenizer"]