Implementation:Shiyu coder Kronos KronosTokenizer From Pretrained

Field	Value
implementation_name	KronosTokenizer_From_Pretrained
repo	Shiyu_coder_Kronos
type	API Doc
source_file	model/kronos.py:L13-177
class	KronosTokenizer(nn.Module, PyTorchModelHubMixin)
implements	Principle:Shiyu_coder_Kronos_Tokenizer_Loading
last_updated	2026-02-09 14:00 GMT

Summary

The KronosTokenizer.from_pretrained class method loads a pre-trained VQ-VAE tokenizer from a HuggingFace Hub repository or a local directory path, returning a fully initialized KronosTokenizer nn.Module on CPU.

API Signature

KronosTokenizer.from_pretrained(
    pretrained_model_name_or_path: str,
    **kwargs
) -> KronosTokenizer

Import

from model import KronosTokenizer
# or
from model.kronos import KronosTokenizer

Parameters

from_pretrained Parameters

Parameter	Type	Description
pretrained_model_name_or_path	str	HuggingFace Hub model ID (e.g., "NeoQuasar/Kronos-Tokenizer-base") or local filesystem path to a saved model directory.
**kwargs	dict	Additional keyword arguments passed to the HuggingFace Hub download and model initialization.

__init__ Parameters (loaded from config)

Parameter	Type	Description
d_in	int	Input dimension (number of features per timestep).
d_model	int	Model hidden dimension.
n_heads	int	Number of attention heads.
ff_dim	int	Feed-forward network dimension.
n_enc_layers	int	Number of encoder Transformer layers.
n_dec_layers	int	Number of decoder Transformer layers.
ffn_dropout_p	float	Dropout probability for feed-forward networks.
attn_dropout_p	float	Dropout probability for attention mechanisms.
resid_dropout_p	float	Dropout probability for residual connections.
s1_bits	int	Number of bits for coarse (s1) tokens in BSQuantizer.
s2_bits	int	Number of bits for fine (s2) tokens in BSQuantizer.
beta	float	Beta parameter for BSQuantizer loss.
gamma0	float	Gamma0 parameter for BSQuantizer.
gamma	float	Gamma parameter for BSQuantizer.
zeta	float	Zeta parameter for BSQuantizer.
group_size	int	Group size parameter for BSQuantizer.

Input

• pretrained_model_name_or_path (str): A HuggingFace Hub model ID string (e.g., "NeoQuasar/Kronos-Tokenizer-base") or a local filesystem path to a directory containing the model weights and configuration.

Output

• KronosTokenizer: A fully initialized nn.Module instance on CPU with pre-trained weights loaded, ready for inference or further device placement.

Dependencies

• torch
• huggingface_hub (provides the PyTorchModelHubMixin base class and download utilities)

Architecture

The loaded KronosTokenizer contains:

KronosTokenizer
  +-- embed: nn.Linear(d_in -> d_model)
  +-- encoder: nn.ModuleList[TransformerBlock x (n_enc_layers - 1)]
  +-- quant_embed: nn.Linear(d_model -> codebook_dim)
  +-- tokenizer: BSQuantizer(s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size)
  +-- post_quant_embed_pre: nn.Linear(s1_bits -> d_model)
  +-- post_quant_embed: nn.Linear(codebook_dim -> d_model)
  +-- decoder: nn.ModuleList[TransformerBlock x (n_dec_layers - 1)]
  +-- head: nn.Linear(d_model -> d_in)

Example

from model import KronosTokenizer

# Load from HuggingFace Hub
tokenizer = KronosTokenizer.from_pretrained("NeoQuasar/Kronos-Tokenizer-base")

# Load from a local path
tokenizer = KronosTokenizer.from_pretrained("/path/to/local/tokenizer/")

# Move to GPU for inference
tokenizer = tokenizer.to("cuda:0")
tokenizer.eval()

Source Code Reference

File: model/kronos.py, lines 13-177.

class KronosTokenizer(nn.Module, PyTorchModelHubMixin):
    def __init__(self, d_in, d_model, n_heads, ff_dim, n_enc_layers, n_dec_layers,
                 ffn_dropout_p, attn_dropout_p, resid_dropout_p,
                 s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):
        super().__init__()
        self.d_in = d_in
        self.d_model = d_model
        # ... encoder, decoder, quantizer initialization ...
        self.tokenizer = BSQuantizer(self.s1_bits, self.s2_bits, beta, gamma0, gamma, zeta, group_size)

The from_pretrained method is inherited from PyTorchModelHubMixin and handles downloading the model configuration and weights from HuggingFace Hub, then calling __init__ with the stored configuration parameters.

Notes

• The model is returned on CPU by default. Move it to the desired device before inference.
• The from_pretrained method is inherited from PyTorchModelHubMixin and is not explicitly defined in the class body.
• The tokenizer should be set to .eval() mode before inference to disable dropout.

Environment & Heuristic Links

• Environment:Shiyu_coder_Kronos_PyTorch_CUDA_Environment
• Environment:Shiyu_coder_Kronos_HuggingFace_Hub_Access