File size: 5,607 Bytes
afd767d 65f5f0b afd767d 65f5f0b afd767d 65f5f0b afd767d 65f5f0b afd767d | 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 130 | # 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 should treat every non-pad
# position as a real token.
kwargs.setdefault("pad_token", chr(1))
kwargs.setdefault("padding_side", "left")
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(0, min(n, self.vocab_size - 1))
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=0, 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
|