guidellm.data.tokenizers

@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'."
        ),
    )

class TokenizerRegistry(RegistryMixin[type[DataTokenizer]]):
    @classmethod
    def create(cls, config: DataTokenizerArgs) -> DataTokenizer:
        """
        Factory method to create a DatasetTokenizer instance based on configuration.

        :param config: A DataTokenizerArgs object containing the configuration.
        """
        kind = config.kind
        tokenizer_cls = cls.get_registered_object(kind)

        if tokenizer_cls is None:
            raise ValueError(
                f"DataTokenizer type '{kind}' is not registered."
                f"Available types: {list(cls.registry.keys()) if cls.registry else []}"
            )

        return tokenizer_cls(config)