Rorical/logos-1b-base

1"""HuggingFace tokenizer wrapper around tiktoken for Logos."""2 3from __future__ import annotations4 5import json6from pathlib import Path7from typing import Dict, Iterable, List, Optional, Tuple8 9import tiktoken10from transformers import PreTrainedTokenizer11 12 13class LogosTokenizer(PreTrainedTokenizer):14    model_input_names = ["input_ids", "attention_mask"]15    vocab_files_names: Dict[str, str] = {}16 17    def __init__(18        self,19        encoding_name: str = "cl100k_base",20        errors: str = "replace",21        **kwargs,22    ):23        self.encoding_name = encoding_name24        self.encoding = tiktoken.get_encoding(encoding_name)25        self.errors = errors26        eos = "<|endoftext|>"27        kwargs.setdefault("eos_token", eos)28        kwargs.setdefault("pad_token", eos)29        kwargs.setdefault("unk_token", eos)30        super().__init__(**kwargs)31 32    @property33    def vocab_size(self) -> int:34        return int(self.encoding.n_vocab)35 36    def get_vocab(self) -> Dict[str, int]:37        return {str(i): i for i in range(self.vocab_size)}38 39    def __len__(self) -> int:40        return self.vocab_size41 42    def _tokenize(self, text: str, **kwargs) -> List[str]:43        ids = self.encoding.encode(44            text,45            allowed_special=kwargs.get("allowed_special", set()),46            disallowed_special=kwargs.get("disallowed_special", ()),47        )48        return [str(i) for i in ids]49 50    def _convert_token_to_id(self, token: str) -> int:51        if token in {self.eos_token, self.pad_token, self.unk_token}:52            return int(self.encoding.eot_token)53        try:54            return int(token)55        except (TypeError, ValueError):56            return int(self.encoding.eot_token)57 58    def _convert_id_to_token(self, index: int) -> str:59        if int(index) == int(self.encoding.eot_token):60            return self.eos_token61        return str(int(index))62 63    def convert_tokens_to_ids(self, tokens):64        if tokens is None:65            return None66        if isinstance(tokens, (list, tuple)):67            return [self._convert_token_to_id(tok) for tok in tokens]68        return self._convert_token_to_id(tokens)69 70    def convert_ids_to_tokens(self, ids, skip_special_tokens: bool = False):71        if ids is None:72            return None73        if isinstance(ids, (list, tuple)):74            return [self.convert_ids_to_tokens(i, skip_special_tokens=skip_special_tokens) for i in ids]75        idx = int(ids)76        if skip_special_tokens and idx == int(self.encoding.eot_token):77            return None78        return self._convert_id_to_token(idx)79 80    def convert_tokens_to_string(self, tokens: Iterable[str]) -> str:81        ids = [self._convert_token_to_id(tok) for tok in tokens]82        return self.encoding.decode(ids, errors=self.errors)83 84    def build_inputs_with_special_tokens(85        self,86        token_ids_0: List[int],87        token_ids_1: Optional[List[int]] = None,88    ) -> List[int]:89        if token_ids_1 is None:90            return list(token_ids_0)91        return list(token_ids_0) + list(token_ids_1)92 93    def get_special_tokens_mask(94        self,95        token_ids_0: List[int],96        token_ids_1: Optional[List[int]] = None,97        already_has_special_tokens: bool = False,98    ) -> List[int]:99        if already_has_special_tokens:100            ids = token_ids_0101        elif token_ids_1 is None:102            ids = token_ids_0103        else:104            ids = token_ids_0 + token_ids_1105        eos_id = int(self.encoding.eot_token)106        return [1 if int(tok) == eos_id else 0 for tok in ids]107 108    def _decode(109        self,110        token_ids: List[int],111        skip_special_tokens: bool = False,112        clean_up_tokenization_spaces: Optional[bool] = None,113        **kwargs,114    ) -> str:115        ids = [int(i) for i in token_ids]116        if skip_special_tokens:117            eos_id = int(self.encoding.eot_token)118            ids = [i for i in ids if i != eos_id]119        return self.encoding.decode(ids, errors=self.errors)120 121    def save_vocabulary(122        self,123        save_directory: str,124        filename_prefix: Optional[str] = None,125    ) -> Tuple[str, ...]:126        path = Path(save_directory)127        path.mkdir(parents=True, exist_ok=True)128        name = f"{filename_prefix + '-' if filename_prefix else ''}logos_tokenizer.json"129        out = path / name130        out.write_text(json.dumps({"encoding_name": self.encoding_name}, indent=2))131        return (str(out),)132 133 134__all__ = ["LogosTokenizer"]135