File size: 2,787 Bytes
7cf5414 | 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 | from typing import Iterable
import torch
from transformers import AutoTokenizer, PreTrainedTokenizerFast
class MagicBERTTokenizer(PreTrainedTokenizerFast):
def card_token_ids(self) -> list[int]:
specials = set(self.all_special_tokens)
vocab = self.get_vocab()
return sorted(token_id for token, token_id in vocab.items() if token not in specials)
def is_card_token(self, token: str) -> bool:
return token not in set(self.all_special_tokens)
def is_card_id(self, token_id: int) -> bool:
token: str = self.convert_ids_to_tokens(token_id) # type: ignore
return bool(token) and self.is_card_token(token)
def convert_card_names_to_ids(self, names: Iterable[str]) -> torch.Tensor:
ids = [self.convert_tokens_to_ids(name) for name in names]
return torch.tensor(ids, dtype=torch.long)
def encode(
self,
text,
text_pair=None,
add_special_tokens=True,
padding="max_length",
truncation=False,
max_length=None,
stride=0,
padding_side=None,
return_tensors=None,
**kwargs,
):
if isinstance(text, list) and text_pair is None:
if not text or isinstance(text[0], str):
return super().encode(
text,
text_pair=text_pair,
add_special_tokens=add_special_tokens,
padding=padding,
truncation=truncation,
max_length=max_length,
stride=stride,
padding_side=padding_side,
return_tensors=return_tensors,
is_split_into_words=True,
**kwargs,
)
if isinstance(text[0], list):
batch_encoding = super().__call__(
text=text,
add_special_tokens=add_special_tokens,
padding=padding,
truncation=truncation,
max_length=max_length,
stride=stride,
padding_side=padding_side,
return_tensors=return_tensors,
is_split_into_words=True,
**kwargs,
)
return batch_encoding["input_ids"]
return super().encode(
text,
text_pair=text_pair,
add_special_tokens=add_special_tokens,
padding=padding,
truncation=truncation,
max_length=max_length,
stride=stride,
padding_side=padding_side,
return_tensors=return_tensors,
**kwargs,
)
MagicBERTTokenizer.register_for_auto_class(AutoTokenizer)
|