Model: Add GenericTokenizer

Because Hugging Face doesnt have a consistant API across their own
libraries
This commit is contained in:
somebody
2023-03-13 17:36:58 -05:00
parent 60793eb121
commit bf8b60ac2d
2 changed files with 26 additions and 2 deletions

23
modeling/tokenizer.py Normal file
View File

@@ -0,0 +1,23 @@
from typing import List, Union
from tokenizers import Tokenizer
import torch
from transformers import PreTrainedTokenizer
class GenericTokenizer:
"""Bridges the gap between Transformers tokenizers and Tokenizers tokenizers. Why they aren't the same, I don't know."""
def __init__(self, tokenizer: Union[Tokenizer, PreTrainedTokenizer]) -> None:
self.tokenizer = tokenizer
# TODO: Get rid of this
self._koboldai_header = []
def encode(self, text: str) -> list:
return self.tokenizer.encode(text).ids
def decode(self, tokens: Union[int, List[int], torch.Tensor]) -> str:
if isinstance(tokens, int):
tokens = [tokens]
return self.tokenizer.decode(tokens)