Upload tokenization_logos.py

tokenization_logos.py

@@ -0,0 +1,134 @@
+"""HuggingFace tokenizer wrapper around tiktoken for Logos."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Tuple
+
+import tiktoken
+from transformers import PreTrainedTokenizer
+
+
+class LogosTokenizer(PreTrainedTokenizer):
+    model_input_names = ["input_ids", "attention_mask"]
+    vocab_files_names: Dict[str, str] = {}
+
+    def __init__(
+        self,
+        encoding_name: str = "cl100k_base",
+        errors: str = "replace",
+        **kwargs,
+    ):
+        self.encoding_name = encoding_name
+        self.encoding = tiktoken.get_encoding(encoding_name)
+        self.errors = errors
+        eos = "<|endoftext|>"
+        kwargs.setdefault("eos_token", eos)
+        kwargs.setdefault("pad_token", eos)
+        kwargs.setdefault("unk_token", eos)
+        super().__init__(**kwargs)
+
+    @property
+    def vocab_size(self) -> int:
+        return int(self.encoding.n_vocab)
+
+    def get_vocab(self) -> Dict[str, int]:
+        return {str(i): i for i in range(self.vocab_size)}
+
+    def __len__(self) -> int:
+        return self.vocab_size
+
+    def _tokenize(self, text: str, **kwargs) -> List[str]:
+        ids = self.encoding.encode(
+            text,
+            allowed_special=kwargs.get("allowed_special", set()),
+            disallowed_special=kwargs.get("disallowed_special", ()),
+        )
+        return [str(i) for i in ids]
+
+    def _convert_token_to_id(self, token: str) -> int:
+        if token in {self.eos_token, self.pad_token, self.unk_token}:
+            return int(self.encoding.eot_token)
+        try:
+            return int(token)
+        except (TypeError, ValueError):
+            return int(self.encoding.eot_token)
+
+    def _convert_id_to_token(self, index: int) -> str:
+        if int(index) == int(self.encoding.eot_token):
+            return self.eos_token
+        return str(int(index))
+
+    def convert_tokens_to_ids(self, tokens):
+        if tokens is None:
+            return None
+        if isinstance(tokens, (list, tuple)):
+            return [self._convert_token_to_id(tok) for tok in tokens]
+        return self._convert_token_to_id(tokens)
+
+    def convert_ids_to_tokens(self, ids, skip_special_tokens: bool = False):
+        if ids is None:
+            return None
+        if isinstance(ids, (list, tuple)):
+            return [self.convert_ids_to_tokens(i, skip_special_tokens=skip_special_tokens) for i in ids]
+        idx = int(ids)
+        if skip_special_tokens and idx == int(self.encoding.eot_token):
+            return None
+        return self._convert_id_to_token(idx)
+
+    def convert_tokens_to_string(self, tokens: Iterable[str]) -> str:
+        ids = [self._convert_token_to_id(tok) for tok in tokens]
+        return self.encoding.decode(ids, errors=self.errors)
+
+    def build_inputs_with_special_tokens(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+    ) -> List[int]:
+        if token_ids_1 is None:
+            return list(token_ids_0)
+        return list(token_ids_0) + list(token_ids_1)
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+        already_has_special_tokens: bool = False,
+    ) -> List[int]:
+        if already_has_special_tokens:
+            ids = token_ids_0
+        elif token_ids_1 is None:
+            ids = token_ids_0
+        else:
+            ids = token_ids_0 + token_ids_1
+        eos_id = int(self.encoding.eot_token)
+        return [1 if int(tok) == eos_id else 0 for tok in ids]
+
+    def _decode(
+        self,
+        token_ids: List[int],
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: Optional[bool] = None,
+        **kwargs,
+    ) -> str:
+        ids = [int(i) for i in token_ids]
+        if skip_special_tokens:
+            eos_id = int(self.encoding.eot_token)
+            ids = [i for i in ids if i != eos_id]
+        return self.encoding.decode(ids, errors=self.errors)
+
+    def save_vocabulary(
+        self,
+        save_directory: str,
+        filename_prefix: Optional[str] = None,
+    ) -> Tuple[str, ...]:
+        path = Path(save_directory)
+        path.mkdir(parents=True, exist_ok=True)
+        name = f"{filename_prefix + '-' if filename_prefix else ''}logos_tokenizer.json"
+        out = path / name
+        out.write_text(json.dumps({"encoding_name": self.encoding_name}, indent=2))
+        return (str(out),)
+
+
+__all__ = ["LogosTokenizer"]