wehub-resource-sync
187 lines
8.5 KiB
Python
187 lines
8.5 KiB
Python
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
|
|
|
|
import json
|
|
import warnings
|
|
from collections.abc import Iterable, Iterator
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from litgpt.utils import fix_and_load_json
|
|
|
|
|
|
class Tokenizer:
|
|
def __init__(self, checkpoint_dir: Path | str) -> None:
|
|
checkpoint_dir = Path(checkpoint_dir)
|
|
if not checkpoint_dir.exists():
|
|
raise NotADirectoryError(f"The checkpoint directory does not exist: {str(checkpoint_dir)}")
|
|
|
|
self.model_name = checkpoint_dir.stem
|
|
self.use_bos = self.check_if_bos_token_used(checkpoint_dir)
|
|
self.bos_id = None
|
|
self.eos_id = None
|
|
|
|
# some checkpoints have both files, `.json` takes precedence
|
|
if (vocabulary_path := checkpoint_dir / "tokenizer.json").is_file():
|
|
from tokenizers import Tokenizer as HFTokenizer
|
|
|
|
self.processor = HFTokenizer.from_file(str(vocabulary_path))
|
|
self.backend = "huggingface"
|
|
|
|
if (special_tokens_path := checkpoint_dir / "tokenizer_config.json").is_file():
|
|
with open(special_tokens_path, encoding="utf-8") as fp:
|
|
config = json.load(fp)
|
|
bos_token = config.get("bos_token")
|
|
eos_token = config.get("eos_token")
|
|
if bos_token is not None and isinstance(bos_token, dict):
|
|
bos_token = bos_token.get("content")
|
|
if eos_token is not None and isinstance(eos_token, dict):
|
|
eos_token = eos_token.get("content")
|
|
self.bos_id = self.token_to_id(bos_token) if bos_token is not None else None
|
|
self.eos_id = self.token_to_id(eos_token) if eos_token is not None else None
|
|
if (special_tokens_path := checkpoint_dir / "generation_config.json").is_file():
|
|
try:
|
|
with open(special_tokens_path, encoding="utf-8") as fp:
|
|
config = json.load(fp)
|
|
except json.JSONDecodeError: # Some files like the Llama 3.2 one have bugs
|
|
warnings.warn(
|
|
f"generation_config.json in '{checkpoint_dir}' contains invalid JSON. "
|
|
"Attempting automatic fix. Verify this file is not corrupted.",
|
|
stacklevel=2,
|
|
)
|
|
with open(special_tokens_path, encoding="utf-8") as fp:
|
|
json_string = fp.read()
|
|
config = fix_and_load_json(json_string)
|
|
if self.bos_id is None:
|
|
self.bos_id = config.get("bos_token_id")
|
|
if self.eos_id is None:
|
|
self.eos_id = config.get("eos_token_id")
|
|
|
|
elif (vocabulary_path := checkpoint_dir / "tokenizer.model").is_file():
|
|
from sentencepiece import SentencePieceProcessor
|
|
|
|
self.processor = SentencePieceProcessor(model_file=str(vocabulary_path))
|
|
self.backend = "sentencepiece"
|
|
self.bos_id = self.processor.bos_id()
|
|
self.eos_id = self.processor.eos_id()
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
# NOTE: A temporary fix until it's resolved on Tokenizers side.
|
|
# LlaMA tokenizer strips leading spaces if to decode a single token at a time.
|
|
# https://github.com/huggingface/transformers/issues/31643
|
|
self.apply_decoding_fix = None
|
|
if (config_path := checkpoint_dir / "tokenizer_config.json").is_file():
|
|
with open(config_path, encoding="utf-8") as fp:
|
|
self.apply_decoding_fix = "LlamaTokenizer" in json.load(fp)["tokenizer_class"]
|
|
|
|
@property
|
|
def vocab_size(self) -> int:
|
|
if self.backend == "huggingface":
|
|
return self.processor.get_vocab_size(with_added_tokens=False)
|
|
if self.backend == "sentencepiece":
|
|
return self.processor.vocab_size()
|
|
raise RuntimeError
|
|
|
|
def token_to_id(self, token: str) -> int:
|
|
if self.backend == "huggingface":
|
|
id_ = self.processor.token_to_id(token)
|
|
elif self.backend == "sentencepiece":
|
|
id_ = self.processor.piece_to_id(token)
|
|
else:
|
|
raise RuntimeError
|
|
if id_ is None:
|
|
raise ValueError(f"token {token!r} not found in the collection.")
|
|
return id_
|
|
|
|
def check_if_bos_token_used(self, checkpoint_dir: Path) -> bool:
|
|
if not (tokenizer_config_path := checkpoint_dir / "tokenizer_config.json").is_file():
|
|
return False
|
|
with open(tokenizer_config_path, encoding="utf-8") as fp:
|
|
config = json.load(fp)
|
|
# for LlaMA-3 tokenizer there is no `add_bos_token` at all and `tokenizer_class` is only
|
|
# `PreTrainedTokenizerFast`
|
|
if checkpoint_dir.stem.startswith(("Meta-Llama-3", "Llama-3")):
|
|
return True
|
|
if checkpoint_dir.stem.startswith("SmolLM2") and checkpoint_dir.name.endswith("Instruct"):
|
|
return True
|
|
if "add_bos_token" in config:
|
|
return config["add_bos_token"]
|
|
# if `add_bos_token` isn't in the config file, but LLaMA tokenizer is used - return True.
|
|
# ex: https://huggingface.co/stabilityai/StableBeluga2/blob/main/tokenizer_config.json#L2
|
|
return config.get("tokenizer_class") == "LlamaTokenizer"
|
|
|
|
def encode(
|
|
self,
|
|
string: str,
|
|
device: torch.device | None = None,
|
|
bos: bool | None = None,
|
|
eos: bool = False,
|
|
max_length: int = -1,
|
|
) -> torch.Tensor:
|
|
if self.backend == "huggingface":
|
|
tokens = self.processor.encode(string).ids
|
|
elif self.backend == "sentencepiece":
|
|
tokens = self.processor.encode(string)
|
|
else:
|
|
raise RuntimeError(f"`{self.backend}` is not supported.")
|
|
if tokens is None:
|
|
raise ValueError("`self.processor` returned tokens of None value.")
|
|
|
|
if bos or (bos is None and self.use_bos):
|
|
if self.bos_id is None:
|
|
raise NotImplementedError("This tokenizer does not have a defined bos token.")
|
|
if not tokens or tokens[0] != self.bos_id:
|
|
tokens = [self.bos_id] + tokens
|
|
# if the processor misbehaves and adds `bos` token no matter what
|
|
elif tokens and tokens[0] == self.bos_id:
|
|
tokens = tokens[1:]
|
|
|
|
if eos and (not tokens or tokens[-1] != self.eos_id):
|
|
tokens = tokens + [self.eos_id]
|
|
# if the processor misbehaves and adds `eos` token no matter what
|
|
elif tokens and tokens[-1] == self.eos_id:
|
|
tokens = tokens[:-1]
|
|
|
|
if max_length > 0:
|
|
tokens = tokens[:max_length]
|
|
return torch.tensor(tokens, dtype=torch.int, device=device)
|
|
|
|
def decode(self, tensor: torch.Tensor) -> str:
|
|
tokens = [tensor.item()] if tensor.ndim == 0 else tensor.tolist()
|
|
if len(tokens) == 1 and self.apply_decoding_fix:
|
|
dummy_token_id = 33 # \x1e
|
|
dummy_token = self.processor.decode([dummy_token_id])
|
|
if dummy_token != "\x1e":
|
|
dummy_token_id = 165 # \x1e is different in salamandra tokenizers
|
|
dummy_token = self.processor.decode([dummy_token_id])
|
|
return self.processor.decode([dummy_token_id] + tokens)[len(dummy_token) :]
|
|
return self.processor.decode(tokens)
|
|
|
|
def decode_stream(self, token_stream: Iterable[torch.Tensor], device: torch.device | None = None) -> Iterator[str]:
|
|
if self.backend == "huggingface":
|
|
try:
|
|
for token in token_stream:
|
|
yield self.decode(token)
|
|
except KeyboardInterrupt:
|
|
return
|
|
elif self.backend == "sentencepiece":
|
|
# TODO: Is there a way to not have to do this?
|
|
# This may actually affect our tokens per second.
|
|
|
|
# sentencepiece does not support decoding token-by-token because it adds spaces based on the surrounding tokens
|
|
# meaning that we need to decode everything each time
|
|
so_far = torch.tensor([], dtype=torch.long, device=device)
|
|
decoded_so_far = ""
|
|
try:
|
|
for token in token_stream:
|
|
so_far = so_far.to(device=token.device)
|
|
so_far = torch.cat((so_far, token.view(-1)))
|
|
decoded_new = self.decode(so_far)
|
|
yield decoded_new[len(decoded_so_far) :]
|
|
decoded_so_far = decoded_new
|
|
except KeyboardInterrupt:
|
|
return
|
|
else:
|
|
raise NotImplementedError(self.backend)
|