gems

Prompt

"""Tokenizer API.""" from __future__ import annotations import dataclasses import enum import functools import typing from typing import ClassVar import dialog import einops from etils import enp from etils import epath from etils import epy from fantasy.gm.utils import _file_cache import jax import jax.numpy as jnp from kauldron.utils import immutabledict import numpy as np from sentencepiece import sentencepiece_model_pb2 import sentencepiece as spm with epy.lazy_imports(): from plotly import graph_objects as go import plotly.express as px _WHITESPACE_CHAR = '▁' class _DisplayEnumType(enum.EnumType): def __repr__(cls): return epy.Lines.make_block( header=cls.__name__, content={value.name: value.value for value in cls}, ) class SpecialTokens(enum.IntEnum, metaclass=_DisplayEnumType): PAD: ClassVar[int] EOS: ClassVar[int] BOS: ClassVar[int] UNK: ClassVar[int] MASK: ClassVar[int] CUSTOM: ClassVar[int] START_OF_TURN: ClassVar[int] END_OF_TURN: ClassVar[int] IMAGE_PLACEHOLDER: ClassVar[int] START_OF_IMAGE: ClassVar[int] END_OF_IMAGE: ClassVar[int] AUDIO_PLACEHOLDER: ClassVar[int] START_OF_AUDIO: ClassVar[int] END_OF_AUDIO: ClassVar[int] BEGIN_OF_TOOL_RESPONSE: ClassVar[int] END_OF_TOOL_RESPONSE: ClassVar[int] class _Fantasy2SpecialTokens(SpecialTokens, enum.IntEnum): PAD = 0 EOS = 1 BOS = 2 UNK = 3 MASK = 4 CUSTOM = 7 START_OF_TURN = 106 END_OF_TURN = 107 class _Fantasy3SpecialTokens(SpecialTokens, enum.IntEnum): PAD = 0 EOS = 1 BOS = 2 UNK = 3 MASK = 4 CUSTOM = 6 START_OF_TURN = 105 END_OF_TURN = 106 IMAGE_PLACEHOLDER = 255999 START_OF_IMAGE = 255999 END_OF_IMAGE = 256000 BEGIN_OF_TOOL_RESPONSE = 50 class _Fantasy4SpecialTokens(SpecialTokens, enum.IntEnum): PAD = 0 EOS = 1 BOS = 2 UNK = 3 MASK = 4 START_OF_TURN = 105 END_OF_TURN = 106 IMAGE_PLACEHOLDER = 258880 START_OF_IMAGE = 255999 END_OF_IMAGE = 258882 AUDIO_PLACEHOLDER = 258881 START_OF_AUDIO = 256000 END_OF_AUDIO = 258883 BEGIN_OF_TOOL_RESPONSE = 50 @dataclasses.dataclass(frozen=True, kw_only=True) class Tokenizer: path: epath.PathLike custom_tokens: dict[int, str] = dataclasses.field(default_factory=dict) VERSION: ClassVar[int | str] = 0 FORBIDDEN_TOKENS: ClassVar[tuple[int, ...]] = () FORMAT: ClassVar[dialog.Format] = dialog.Format.FANTASY3 FORMAT_TO_CONVERT: ClassVar[dialog.Format | None] = None def __post_init__(self): immutabledict.freeze_dict_attrs(self, ['custom_tokens']) @classmethod def from_version(cls, version: int | str) -> Tokenizer: if version == 2: return Fantasy2Tokenizer() elif version == 3: return Fantasy3Tokenizer() elif version == '3n': return Fantasy3nTokenizer() elif version == 4: return Fantasy4Tokenizer() else: raise ValueError(f'Unsupported tokenizer version: {version}') def encode( self, text: str | list[str], *, add_bos: bool = False, add_eos: bool = False, ) -> list[int]: if isinstance(text, str): if self.FORMAT_TO_CONVERT: text = self.FORMAT_TO_CONVERT.from_fantasy4(text) token_ids = self._sp.EncodeAsIds(text) else: text = [t.replace(' ', _WHITESPACE_CHAR) for t in text] if self.FORMAT_TO_CONVERT: text = [self.FORMAT_TO_CONVERT.from_fantasy4(t) for t in text] token_ids = [self._sp.PieceToId(t) for t in text] if self.special_tokens.UNK in token_ids: index = token_ids.index(self.special_tokens.UNK) raise ValueError( f'Cannot tokenize {text!r}. Token {text[index]!r} is an unknown token.' ) if add_bos: token_ids.insert(0, self.special_tokens.BOS) if add_eos: token_ids.append(self.special_tokens.EOS) return token_ids def decode(self, ids: int | list[int] | enp.typing.Array) -> str: if isinstance(ids, int): ids = [ids] elif enp.lazy.is_array(ids): ids = typing.cast(np.ndarray, ids) if ids.ndim == 0: ids = [ids.item()] elif ids.ndim == 1: ids = ids.tolist() else: raise ValueError(f'Array must be 0 or 1 dimensional, got {ids.shape}.') text = self._sp.DecodeIds(ids) if self.FORMAT_TO_CONVERT: text = self.FORMAT_TO_CONVERT.to_fantasy4(text) return text def split(self, text: str) -> list[str]: return [_real_whitespaces(t) for t in self._sp.EncodeAsPieces(text)] @functools.cached_property def vocab_size(self) -> int: return self._sp.GetPieceSize() @functools.cached_property def tokens(self) -> list[str]: return [_real_whitespaces(self._sp.IdToPiece(i)) for i in range(self.vocab_size)] @functools.cached_property def special_tokens(self) -> type[SpecialTokens]: raise NotImplementedError(f'{type(self).__qualname__} does not define special tokens.') @functools.cached_property def _sp(self) -> spm.SentencePieceProcessor: sp = spm.SentencePieceProcessor() model_file_path = _file_cache.maybe_get_from_cache( remote_file_path=self.path, cache_subdir='tokenizer', ) model_proto = epath.Path(model_file_path).read_bytes() if self.custom_tokens: model_proto = self._add_custom_tokens(model_proto) sp.LoadFromSerializedProto(model_proto) return sp def _add_custom_tokens(self, serialized_proto: bytes) -> bytes: proto = sentencepiece_model_pb2.ModelProto() proto.ParseFromString(serialized_proto) for i, token in self.custom_tokens.items(): if i < 0 or i > 98: raise ValueError(f'Custom token id {i} for {token!r} is not in [1, 98].') piece = proto.pieces[self.special_tokens.CUSTOM + i] if piece.piece != f'<unused{i}>': raise AssertionError( f'Expected custom token id {i} for {token!r} to be `<unused{i}>`,' f' but was {piece.piece}. This indicates the voab file isn\'t as expected.' ) piece.piece = token if proto.trainer_spec.user_defined_symbols: for index, symbol in enumerate(proto.trainer_spec.user_defined_symbols): if symbol == f'<unused{i}>': break else: raise AssertionError( f'Expected custom token id {i} for {token!r} to be in user_defined_symbols, but it was not found.' ) proto.trainer_spec.user_defined_symbols[index] = token return proto.SerializeToString() def plot_logits( self, logits: enp.typing.Array, *, keep_top: int = 30, ) -> go.Figure: if logits.ndim == 2 and logits.shape[0] == 1: logits = einops.rearrange(logits, '1 d -> d') if logits.ndim != 1: raise ValueError('`plot_logits` expects logits for a single token, got' f' {logits.shape}.') probs = jax.nn.softmax(logits) indices = jnp.argsort(probs) indices = indices[-keep_top:][::-1] probs = probs[indices].astype(np.float32) words = [repr(self.tokens[i]) for i in indices] fig = px.bar(x=words, y=probs) fig.update_layout( title='Probability Distribution of Tokens', xaxis_title='Tokens', yaxis_title='Probability', ) return fig def __getstate__(self): return {f.name: getattr(self, f.name) for f in dataclasses.fields(self)} def __setstate__(self, state): self.__init__(**state) @dataclasses.dataclass(frozen=True) class Fantasy2Tokenizer(Tokenizer): path: epath.PathLike = 'gs://fantasy-data/tokenizers/tokenizer_fantasy2.model' special_tokens = _Fantasy2SpecialTokens VERSION = 2 @dataclasses.dataclass(frozen=True) class Fantasy3Tokenizer(Tokenizer): path: epath.PathLike = 'gs://fantasy-data/tokenizers/tokenizer_fantasy3.model' special_tokens = _Fantasy3SpecialTokens FORBIDDEN_TOKENS = (special_tokens.START_OF_IMAGE, special_tokens.END_OF_IMAGE) VERSION = 3 @dataclasses.dataclass(frozen=True) class Fantasy3nTokenizer(Tokenizer): path: epath.PathLike = 'gs://fantasy-data/tokenizers/tokenizer_fantasy3n.model' special_tokens = _Fantasy3SpecialTokens FORBIDDEN_TOKENS = (special_tokens.START_OF_IMAGE, special_tokens.END_OF_IMAGE) VERSION = '3n' @dataclasses.dataclass(frozen=True) class Fantasy4Tokenizer(Tokenizer): path: epath.PathLike = 'gs://fantasy-data/tokenizers/tokenizer_fantasy4.model' special_tokens = _Fantasy4SpecialTokens VERSION = 4 FORMAT: ClassVar[dialog.Format] = dialog.Format.FANTASY4 def _real_whitespaces(text: str) -> str: return text.replace(_WHITESPACE_CHAR, ' ') I want to perform an exhaustive, line-by-line static analysis of the provided code to identify every single error, mock, dummy, stub, placeholder, and hidden logical flaw, so that the final output is a 100% complete, verified list of real issues with absolutely zero omissions or hallucinations. CRITICAL CONSTRAINTS (DO NOT BREAK THEM): 1. Read every single character from start to finish. Do not skip, summarize, or abbreviate any part of the code. 2. Identify ALL structural and logical flaws: mocks, dummies, stubs, placeholders, syntax errors, and hidden runtime exceptions. 3. Theoretically execute the code paths to uncover non-obvious errors that would occur in practice. 4. Focus EXCLUSIVELY on real, verifiable errors. Do not invent, hallucinate, or assume errors that do not exist. Write down exactly what you find, and nothing more. 5. NO polite filler, NO introductions, NO summaries, NO explanations outside the requested format. OUTPUT FORMAT: Provide the output strictly in the following structure: [ERROR LIST] - Line [X]: [Exact error description] ... [END OF LIST] FILE CLOSED. ALL ERRORS LISTED.