guidellm.data.tokenizers.huggingface

@TokenizerRegistry.register(["huggingface_auto", "hf_auto"])
class HuggingFaceTokenizer(DataTokenizer):
    """Tokenizer for Hugging Face models."""

    def __init__(
        self,
        config: HuggingFaceTokenizerArgs,
    ) -> None:
        if config.model is None:
            raise ValueError("The 'name' field must be provided")

        self._config = config
        self._tokenizer: None | PreTrainedTokenizerBase = None

    def __call__(self) -> PreTrainedTokenizerBase:
        if self._tokenizer is not None:
            return self._tokenizer
        else:
            from_pretrained = AutoTokenizer.from_pretrained(
                self._config.model,
                **self._config.load_kwargs,
            )
            self._tokenizer = from_pretrained
            return from_pretrained

@DataTokenizerArgs.register(["huggingface_auto", "hf_auto"])
class HuggingFaceTokenizerArgs(DataTokenizerArgs):
    """Model for Hugging Face tokenizer arguments."""

    kind: Literal["huggingface_auto", "hf_auto"] = Field(
        default="huggingface_auto",
        description="Type identifier for the HuggingFace tokenizer.",
    )
    load_kwargs: dict[str, Any] = Field(
        default_factory=dict,
        examples=[{"use_fast": True, "revision": "main"}],
        description=(
            "Optional additional arguments to pass to the HuggingFace tokenizer's "
            "from_pretrained method, such as 'use_fast' or 'revision'."
        ),
    )