| """HydrAMP amino-acid tokenizer.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| from typing import Any |
|
|
| from transformers import PreTrainedTokenizer |
|
|
| DEFAULT_TOKENS = ["<pad>"] + list("ACDEFGHIKLMNPQRSTVWY") |
|
|
|
|
| class HydrAMPAATokenizer(PreTrainedTokenizer): |
| """Character-level amino-acid tokenizer for HydrAMP.""" |
|
|
| vocab_files_names = {"vocab_file": "vocab.json"} |
| model_input_names = ["input_ids", "attention_mask"] |
|
|
| def __init__(self, vocab_file: str | None = None, **kwargs: Any) -> None: |
| if vocab_file is not None: |
| payload = json.loads(Path(vocab_file).read_text()) |
| self._token_to_id = {str(k): int(v) for k, v in payload.items()} |
| else: |
| self._token_to_id = {token: idx for idx, token in enumerate(DEFAULT_TOKENS)} |
| self._id_to_token = {idx: token for token, idx in self._token_to_id.items()} |
| self._strict_unknown = bool(kwargs.pop("strict_unknown", True)) |
|
|
| kwargs.setdefault("pad_token", "<pad>") |
| kwargs.setdefault("model_max_length", 25) |
| kwargs.setdefault("padding_side", "right") |
| super().__init__(**kwargs) |
|
|
| if self.pad_token not in self._token_to_id: |
| raise ValueError(f"pad token '{self.pad_token}' not present in tokenizer vocabulary.") |
| self.pad_token_id = self._token_to_id[self.pad_token] |
|
|
| @property |
| def vocab_size(self) -> int: |
| return len(self._token_to_id) |
|
|
| def get_vocab(self) -> dict[str, int]: |
| return dict(self._token_to_id) |
|
|
| def _tokenize(self, text: str) -> list[str]: |
| return list(text.strip().upper()) |
|
|
| def _convert_token_to_id(self, token: str) -> int: |
| if token in self._token_to_id: |
| return self._token_to_id[token] |
| if self._strict_unknown: |
| raise ValueError(f"Unknown amino-acid token '{token}'.") |
| return self.pad_token_id |
|
|
| def _convert_id_to_token(self, index: int) -> str: |
| return self._id_to_token.get(index, self.pad_token) |
|
|
| def convert_tokens_to_string(self, tokens: list[str]) -> str: |
| return "".join(token for token in tokens if token != self.pad_token) |
|
|
| def build_inputs_with_special_tokens(self, token_ids_0: list[int], token_ids_1: list[int] | None = None) -> list[int]: |
| if token_ids_1 is not None: |
| return token_ids_0 + token_ids_1 |
| return token_ids_0 |
|
|
| def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]: |
| save_dir = Path(save_directory) |
| save_dir.mkdir(parents=True, exist_ok=True) |
| filename = "vocab.json" if filename_prefix is None else f"{filename_prefix}-vocab.json" |
| vocab_path = save_dir / filename |
| vocab_path.write_text(json.dumps(self._token_to_id, indent=2, sort_keys=True) + "\n") |
| return (str(vocab_path),) |
|
|
|
|
