File size: 5,578 Bytes
9e26fe9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | # Based on togethercomputer/evo-1-131k-base/tokenizer.py (Apache 2.0).
# Adapted: minor fix (typing), Dict -> dict for forward-compatibility.
"""ByteTokenizer for Evo1.
Maps text to/from raw UTF-8 byte values, vocab_size = 512 (no real vocab,
just the byte range padded out).
"""
from __future__ import annotations
from os import PathLike
from typing import List, Tuple
import numpy as np
from transformers.tokenization_utils import PreTrainedTokenizer
from transformers.tokenization_utils_base import BatchEncoding, TruncationStrategy
from transformers.utils.generic import PaddingStrategy
EMPTY: str = ""
class ByteTokenizer(PreTrainedTokenizer):
"""UTF-8 byte-level tokenizer (vocab_size = 512).
Special tokens follow the original ``CharLevelTokenizer`` convention:
* ``\x00`` (chr(0)) is end-of-document / end-of-sequence.
* ``\x01`` (chr(1)) is the padding token.
These are wired up so ``HF`` tokenization helpers (``padding=True``,
``model_max_length``, etc.) work with this tokenizer.
"""
def __init__(self, byte_level: bool = True, **kwargs):
# Defaults; ``tokenizer_config.json`` may override via kwargs.
# Padding token: byte 0x01 (matches the original CharLevelTokenizer
# pad convention). We deliberately do NOT set eos_token / cls_token:
# Evo1 is a pure byte-level model with no special tokens added at
# encoding time, so downstream pooling logic (e.g. mRNABench) should
# treat every non-pad position as a real token.
kwargs.setdefault("pad_token", chr(1))
super().__init__(byte_level=byte_level, **kwargs)
# The model only consumes input_ids and attention_mask (no segment ids).
self.model_input_names = ["input_ids", "attention_mask"]
@property
def vocab_size(self) -> int:
return 512
@property
def byte_level(self) -> bool:
return self.init_kwargs.get("byte_level", True)
def get_vocab(self) -> dict:
return {chr(i): i for i in range(self.vocab_size)}
def __len__(self) -> int:
return self.vocab_size
def clamp(self, n: int) -> int:
return max(32, min(n, self.vocab_size))
def _tokenize(self, text: str, **kwargs) -> List[str]:
return list(text)
def byte_tokenize(self, text: str) -> np.ndarray:
return np.frombuffer(text.encode("utf-8"), dtype=np.uint8)
def _convert_token_to_id(self, token: str) -> int:
return self.clamp(ord(token))
def _convert_id_to_token(self, index: int) -> str:
return chr(self.clamp(index))
def convert_tokens_to_string(self, tokens: List[str]) -> str:
return EMPTY.join(tokens)
def _decode(self, token_ids: List[int], **kwargs) -> str:
# numpy >= 2 disallows clipping a uint8 array to a value outside [0, 255]
# (vocab_size=512). Convert to int16 first, clip, then back to uint8.
indices = np.asarray(token_ids, dtype=np.int16)
indices = indices.clip(min=32, max=255).astype(np.uint8)
return indices.tobytes().decode("utf-8", errors="replace")
def _encode_plus(self, text: str, **kwargs) -> BatchEncoding:
first_ids = self.byte_tokenize(text).tolist()
return self.prepare_for_model(
first_ids,
pair_ids=None,
add_special_tokens=kwargs.get("add_special_tokens", False),
padding=kwargs.get("padding_strategy", PaddingStrategy.DO_NOT_PAD).value,
truncation=kwargs.get("truncation_strategy", TruncationStrategy.DO_NOT_TRUNCATE).value,
max_length=kwargs.get("max_length"),
stride=kwargs.get("stride", 0),
pad_to_multiple_of=kwargs.get("pad_to_multiple_of"),
return_tensors=kwargs.get("return_tensors"),
prepend_batch_axis=True,
return_attention_mask=kwargs.get("return_attention_mask"),
return_token_type_ids=kwargs.get("return_token_type_ids"),
return_overflowing_tokens=kwargs.get("return_overflowing_tokens", False),
return_special_tokens_mask=kwargs.get("return_special_tokens_mask", False),
return_length=kwargs.get("return_length", False),
verbose=kwargs.get("verbose", True),
)
def _batch_encode_plus(self, batch_text_or_text_pairs, **kwargs) -> BatchEncoding:
input_ids = [(self.byte_tokenize(t).tolist(), None) for t in batch_text_or_text_pairs]
return self._batch_prepare_for_model(
input_ids,
add_special_tokens=kwargs.get("add_special_tokens", False),
padding_strategy=kwargs.get("padding_strategy", PaddingStrategy.DO_NOT_PAD),
truncation_strategy=kwargs.get("truncation_strategy", TruncationStrategy.DO_NOT_TRUNCATE),
max_length=kwargs.get("max_length"),
stride=kwargs.get("stride", 0),
pad_to_multiple_of=kwargs.get("pad_to_multiple_of"),
return_attention_mask=kwargs.get("return_attention_mask"),
return_token_type_ids=kwargs.get("return_token_type_ids"),
return_overflowing_tokens=kwargs.get("return_overflowing_tokens", False),
return_special_tokens_mask=kwargs.get("return_special_tokens_mask", False),
return_length=kwargs.get("return_length", False),
return_tensors=kwargs.get("return_tensors"),
verbose=kwargs.get("verbose", True),
)
def _save_pretrained(
self, save_directory: str | PathLike, file_names: Tuple[str], **kwargs
) -> Tuple[str]:
return file_names
|