hf tokenizer support added

Author: Ilia Kulikov

src/fairseq2/data/text/tokenizers/huggingface_tokenizer.py

@@ -0,0 +1,144 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import final
+
+import torch
+from torch import Tensor
+from typing_extensions import override
+
+from fairseq2.data import VocabularyInfo
+from fairseq2.data.text.tokenizers import (
+    TextTokenDecoder,
+    TextTokenEncoder,
+)
+from fairseq2.typing import Device
+from transformers import AutoTokenizer
+
+
+@final
+class HuggingfaceTokenizerEncoder(TextTokenEncoder):
+    """Represents a tiktoken decoder."""
+
+    _tokenizer: AutoTokenizer
+    _prefix_indices: list[int]
+    _suffix_indices: list[int]
+    _prefix_index_tensor: Tensor | None
+    _suffix_index_tensor: Tensor | None
+    _device: Device | None
+    _pin_memory: bool
+
+    def __init__(
+        self,
+        tokenizer: AutoTokenizer,
+        *,
+        prefix_tokens: Sequence[str] | None = None,
+        suffix_tokens: Sequence[str] | None = None,
+        device: Device | None = None,
+        pin_memory: bool = False,
+    ) -> None:
+        """
+        :param tokenizer:
+            The huggingface :class:`AutoTokenizer` object.
+        :param prefix_tokens:
+            The prefix tokens to encode with input text.
+        :param suffix_tokens:
+            The suffix tokens to encode with input text.
+        :param device:
+            The device on which to construct tensors.
+        :param pin_memory:
+            If ``True``, uses pinned memory while constructing tensors.
+        """
+        self._tokenizer = tokenizer
+
+        # Prefix
+        if prefix_tokens:
+            self._prefix_indices = self._tokenizer.convert_tokens_to_ids(prefix_tokens)
+
+            self._prefix_index_tensor = torch.tensor(
+                self._prefix_indices, dtype=torch.int64, device=device
+            )
+        else:
+            self._prefix_indices = []
+
+            self._prefix_index_tensor = None
+
+        # Suffix
+        if suffix_tokens:
+            self._suffix_indices = self._tokenizer.convert_tokens_to_ids(suffix_tokens)
+
+            self._suffix_index_tensor = torch.tensor(
+                self._suffix_indices, dtype=torch.int64, device=device
+            )
+        else:
+            self._suffix_indices = []
+
+            self._suffix_index_tensor = None
+
+        self._device = device
+        self._pin_memory = pin_memory
+
+    @override
+    def __call__(self, text: str) -> Tensor:
+        # fairseq2 tokenizer adds special tokens on its own
+        indices = self._tokenizer.encode(text, add_special_tokens=False)
+
+        if self._prefix_indices:
+            indices = self._prefix_indices + indices
+
+        if self._suffix_indices:
+            indices.extend(self._suffix_indices)
+
+        return torch.tensor(
+            indices, dtype=torch.int64, device=self._device, pin_memory=self._pin_memory
+        )
+
+    @override
+    def encode_as_tokens(self, text: str) -> list[str]:
+        indices = self(text).tolist()
+
+        tokens = self._tokenizer.convert_tds_to_tokens(indices)
+
+        return tokens
+
+    @property
+    @override
+    def prefix_indices(self) -> Tensor | None:
+        return self._prefix_index_tensor
+
+    @property
+    @override
+    def suffix_indices(self) -> Tensor | None:
+        return self._suffix_index_tensor
+
+
+@final
+class HuggingfaceTokenizerDecoder(TextTokenDecoder):
+    """Represents a tiktoken decoder."""
+
+    _tokenizer: AutoTokenizer
+
+    def __init__(self, tokenizer: AutoTokenizer) -> None:
+        self._tokenizer = tokenizer
+
+    @override
+    def __call__(self, token_indices: Tensor) -> str:
+        if token_indices.dim() != 1:
+            raise ValueError(
+                f"`token_indices` must be one dimensional, but has {token_indices.dim()} dimensions instead."
+            )
+
+        return self._tokenizer.decode(token_indices)
+
+    @override
+    def decode_from_tokens(self, tokens: Sequence[str]) -> str:
+        indices = self._tokenizer.convert_tokens_to_ids(tokens)
+
+        return self._tokenizer.decode(indices)

src/fairseq2/data/text/tokenizers/llama.py

@@ -24,7 +24,97 @@
     TiktokenEncoder,
     TiktokenModel,
 )
+from fairseq2.data.text.tokenizers.huggingface_tokenizer import (
+    HuggingfaceTokenizerEncoder,
+    HuggingfaceTokenizerDecoder,
+)
 from fairseq2.typing import Device
+from transformers import AutoTokenizer
+
+
+@final
+class LLaMA3TokenizerHuggingFace(TextTokenizer):
+    """Represents a HuggingFace version of LLama 3 tokenizer"""
+
+    _tokenizer: AutoTokenizer
+    _bos_token: str
+    _eos_token: str
+
+    def __init__(self, path: Path) -> None:
+
+        self._tokenizer = AutoTokenizer.from_pretrained(path)
+
+        self._eos_token = self._tokenizer.special_tokens_map["eos_token"]
+        self._bos_token = self._tokenizer.special_tokens_map["bos_token"]
+
+    @override
+    def create_encoder(
+        self,
+        *,
+        task: str | None = None,
+        lang: str | None = None,
+        mode: str | None = None,
+        device: Device | None = None,
+        pin_memory: bool = False,
+    ) -> TiktokenEncoder:
+        if task is not None:
+            raise ValueError(f"`task` must be `None`, but is '{task}' instead.")
+
+        if lang is not None:
+            raise ValueError(f"`lang` must be `None`, but is '{lang}' instead.")
+
+        match mode:
+            case None | "default":
+                prefix_tokens = [self._bos_token]
+                suffix_tokens = [self._eos_token]
+            case "prompt":
+                prefix_tokens = [self._bos_token]
+                # In prompt mode, we expect the generator to finish the sequence.
+                suffix_tokens = []
+            case "prompt_response":
+                prefix_tokens = []
+                suffix_tokens = [self._eos_token]
+            case "as_is":
+                prefix_tokens = []
+                suffix_tokens = []
+            case _:
+                raise ValueError(
+                    f"`mode` must be one of the following values, but is '{mode}' instead: default, prompt, prompt_response, as_is"
+                )
+
+        return HuggingfaceTokenizerEncoder(
+            self._tokenizer,
+            prefix_tokens=prefix_tokens,
+            suffix_tokens=suffix_tokens,
+            device=device,
+            pin_memory=pin_memory,
+        )
+
+    @override
+    def create_raw_encoder(
+        self, *, device: Device | None = None, pin_memory: bool = False
+    ) -> TiktokenEncoder:
+        return HuggingfaceTokenizerEncoder(
+            self._tokenizer, device=device, pin_memory=pin_memory
+        )
+
+    @override
+    def create_decoder(self) -> TiktokenDecoder:
+        return HuggingfaceTokenizerDecoder(self._model)
+
+    @property
+    @override
+    def vocab_info(self) -> VocabularyInfo:
+        bos_idx = self._tokenizer.convert_tokens_to_ids(self._bos_token)
+        eos_idx = self._tokenizer.convert_tokens_to_ids(self._eos_token)
+        vocab_info = VocabularyInfo(
+            size=len(self._tokenizer),
+            bos_idx=bos_idx,
+            eos_idx=eos_idx,
+            unk_idx=None,
+            pad_idx=None,
+        )
+        return vocab_info
 
 
 @final
@@ -139,6 +229,27 @@ def vocab_info(self) -> VocabularyInfo:
 
 
 def load_llama_tokenizer(path: Path, card: AssetCard) -> TextTokenizer:
+
+    # first check if this is HuggingFace tokenizer
+    try:
+        use_hf = card.field("use_hf_tokenizer").as_(bool)
+    except AssetCardFieldNotFoundError:
+        use_hf = False
+    except AssetCardError as ex:
+        raise text_tokenizer_asset_card_error(card.name) from ex
+
+    if use_hf:
+        try:
+            return LLaMA3TokenizerHuggingFace(path)
+        except ValueError as ex:
+            raise TextTokenizerLoadError(
+                card.name, f"The '{card.name}' asset card does not contain a valid text tokenizer configuration of the '{LLAMA_TOKENIZER_FAMILY}' family. See the nested exception for details."  # fmt: skip
+            ) from ex
+        except RuntimeError as ex:
+            raise TextTokenizerLoadError(
+                card.name, f"The '{card.name}' text tokenizer cannot be loaded. See the nested exception for details."  # fmt: skip
+            ) from ex
+
     try:
         use_v2 = card.field("use_v2_tokenizer").as_(bool)
     except AssetCardFieldNotFoundError: