transformers 5.0.0__py3-none-any.whl → 5.0.0rc0__py3-none-any.whl

transformers/models/janus/modular_janus.py

@@ -1,3 +1,4 @@
+# coding=utf-8
 # Copyright 2025 Deepseek AI and The HuggingFace Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,13 +16,16 @@
 import copy
 from collections.abc import Callable, Iterable
 from dataclasses import dataclass
+from typing import Optional, Union
 
 import numpy as np
 import torch
 import torch.nn.functional as F
+import torch.utils.checkpoint
 from torch import nn
 
-from ... import initialization as init
+from transformers.models.blip.image_processing_blip import BlipImageProcessor
+
 from ...activations import ACT2FN
 from ...cache_utils import Cache
 from ...configuration_utils import PreTrainedConfig
@@ -41,7 +45,7 @@ from ...image_utils import (
     valid_images,
     validate_preprocess_arguments,
 )
-from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ModelOutput
+from ...modeling_outputs import ModelOutput
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import ImagesKwargs, Unpack
 from ...utils import (
@@ -52,10 +56,8 @@ from ...utils import (
     filter_out_non_signature_kwargs,
     is_vision_available,
     logging,
-    torch_compilable_check,
 )
 from ..auto import CONFIG_MAPPING, AutoConfig, AutoModel
-from ..blip.image_processing_blip import BlipImageProcessor
 from ..blip_2.modeling_blip_2 import Blip2VisionModel
 from ..chameleon.configuration_chameleon import ChameleonVQVAEConfig
 from ..chameleon.modeling_chameleon import (
@@ -389,11 +391,6 @@ class JanusPreTrainedModel(PreTrainedModel):
 
     _can_compile_fullgraph = True
 
-    def _init_weights(self, module):
-        super()._init_weights(module)
-        if isinstance(module, JanusVisionEmbeddings):
-            init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
-
 
 @dataclass
 @auto_docstring(
@@ -409,8 +406,8 @@ class JanusVQVAEOutput(ModelOutput):
         Embedding loss.
     """
 
-    decoded_pixel_values: torch.FloatTensor | None = None
-    embedding_loss: torch.FloatTensor | None = None
+    decoded_pixel_values: Optional[torch.FloatTensor] = None
+    embedding_loss: Optional[torch.FloatTensor] = None
 
 
 class JanusBaseModelOutputWithPast(IdeficsBaseModelOutputWithPast):
@@ -473,7 +470,7 @@ class JanusVisionAttention(nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
+        attention_mask: Optional[torch.Tensor] = None,
         **kwargs: Unpack[TransformersKwargs],
     ):
         batch_size, seq_len, _ = hidden_states.size()
@@ -552,42 +549,10 @@ class JanusVisionEncoder(SiglipEncoder):
 
 
 class JanusVisionModel(Blip2VisionModel):
-    _can_record_outputs = {
-        "hidden_states": JanusVisionEncoderLayer,
-        "attentions": JanusVisionAttention,
-    }
-
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
         self.encoder = JanusVisionEncoder(config)
 
-    def forward(
-        self,
-        pixel_values: torch.FloatTensor | None = None,
-        interpolate_pos_encoding: bool = False,
-        **kwargs: Unpack[TransformersKwargs],
-    ) -> tuple | BaseModelOutputWithPooling:
-        if pixel_values is None:
-            raise ValueError("You have to specify pixel_values")
-
-        hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
-
-        encoder_outputs: BaseModelOutput = self.encoder(
-            inputs_embeds=hidden_states,
-            **kwargs,
-        )
-
-        last_hidden_state = encoder_outputs.last_hidden_state
-        last_hidden_state = self.post_layernorm(last_hidden_state)
-
-        pooled_output = last_hidden_state[:, 0, :]
-        pooled_output = self.post_layernorm(pooled_output)
-
-        return BaseModelOutputWithPooling(
-            last_hidden_state=last_hidden_state,
-            pooler_output=pooled_output,
-        )
-
 
 class JanusVisionAlignerMLP(nn.Module):
     def __init__(self, config: JanusVisionConfig):
@@ -824,10 +789,6 @@ class JanusVQVAE(ChameleonVQVAE):
         "JanusVQVAEResnetBlock",
         "JanusVQVAEVectorQuantizer",
     ]
-    _can_record_outputs = {
-        "hidden_states": JanusVQVAEResnetBlock,
-        "attentions": JanusVQVAEAttnBlock,
-    }
     main_input_name = "pixel_values"
 
     def __init__(self, config: JanusVQVAEConfig):
@@ -862,13 +823,12 @@ class JanusVQVAE(ChameleonVQVAE):
     def forward(
         self,
         pixel_values: torch.FloatTensor,
-        **kwargs,
     ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
         batch_size = pixel_values.shape[0]
-        encode_outputs = self.encode(pixel_values, return_dict=True, **kwargs)
-        decoded_pixel_values = self.decode(encode_outputs.image_tokens.view(batch_size, -1))
+        quant, embedding_loss, indices = self.encode(pixel_values)
+        decoded_pixel_values = self.decode(indices.view(batch_size, -1))
 
-        return JanusVQVAEOutput(decoded_pixel_values, encode_outputs.embedding_loss)
+        return JanusVQVAEOutput(decoded_pixel_values, embedding_loss)
 
 
 class JanusVQVAEAlignerMLP(nn.Module):
@@ -938,15 +898,10 @@ class JanusModel(JanusPreTrainedModel):
     def set_input_embeddings(self, value):
         self.language_model.set_input_embeddings(value)
 
-    @can_return_tuple
-    @auto_docstring
-    def get_image_features(
-        self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]
-    ) -> tuple | BaseModelOutputWithPooling:
-        vision_outputs = self.vision_model(pixel_values, return_dict=True, **kwargs)
-        vision_outputs.pooler_output = self.aligner(vision_outputs.last_hidden_state)
-
-        return vision_outputs
+    def get_image_features(self, pixel_values):
+        image_embeds = self.vision_model(pixel_values)
+        image_embeds = self.aligner(image_embeds.last_hidden_state)
+        return image_embeds
 
     def get_placeholder_mask(
         self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor
@@ -964,29 +919,29 @@ class JanusModel(JanusPreTrainedModel):
             special_image_mask = input_ids == self.config.image_token_id
 
         n_image_tokens = special_image_mask.sum()
-        n_image_features = image_features.shape[0] * image_features.shape[1]
         special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
-        torch_compilable_check(
-            inputs_embeds[special_image_mask].numel() == image_features.numel(),
-            f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}",
-        )
+        if inputs_embeds[special_image_mask].numel() != image_features.numel():
+            n_image_features = image_features.shape[0] * image_features.shape[1]
+            raise ValueError(
+                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+            )
         return special_image_mask
 
     @can_return_tuple
     @auto_docstring
     def forward(
         self,
-        input_ids: torch.LongTensor | None = None,
-        pixel_values: torch.FloatTensor | None = None,
-        attention_mask: torch.Tensor | None = None,
-        position_ids: torch.LongTensor | None = None,
-        past_key_values: Cache | None = None,
-        cache_position: torch.LongTensor | None = None,
-        inputs_embeds: torch.FloatTensor | None = None,
-        use_cache: bool | None = None,
-        logits_to_keep: int | torch.Tensor = 0,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
         **kwargs,
-    ) -> JanusBaseModelOutputWithPast:
+    ):
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError(
                 "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
@@ -995,7 +950,7 @@ class JanusModel(JanusPreTrainedModel):
             inputs_embeds = self.get_input_embeddings()(input_ids)
 
         if pixel_values is not None:
-            image_embeds = self.get_image_features(pixel_values, return_dict=True).pooler_output
+            image_embeds = self.get_image_features(pixel_values)
             image_features = image_embeds.reshape(-1, inputs_embeds.shape[-1])
             image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
             image_attention_mask = self.get_placeholder_mask(
@@ -1052,18 +1007,18 @@ class JanusForConditionalGeneration(JanusPreTrainedModel, GenerationMixin):
     @auto_docstring
     def forward(
         self,
-        input_ids: torch.LongTensor | None = None,
-        pixel_values: torch.FloatTensor | None = None,
-        attention_mask: torch.Tensor | None = None,
-        position_ids: torch.LongTensor | None = None,
-        past_key_values: Cache | None = None,
-        cache_position: torch.LongTensor | None = None,
-        inputs_embeds: torch.FloatTensor | None = None,
-        labels: torch.LongTensor | None = None,
-        use_cache: bool | None = None,
-        logits_to_keep: int | torch.Tensor = 0,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
         **kwargs: Unpack[TransformersKwargs],
-    ) -> JanusCausalLMOutputWithPast:
+    ):
         r"""
         labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
             Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
@@ -1110,7 +1065,6 @@ class JanusForConditionalGeneration(JanusPreTrainedModel, GenerationMixin):
         inputs_embeds=None,
         cache_position=None,
         logits_to_keep=None,
-        is_first_iteration=False,
         **kwargs,
     ):
         # Overwritten -- extra custom processing
@@ -1122,15 +1076,12 @@ class JanusForConditionalGeneration(JanusPreTrainedModel, GenerationMixin):
             attention_mask=attention_mask,
             cache_position=cache_position,
             logits_to_keep=logits_to_keep,
-            is_first_iteration=is_first_iteration,
             **kwargs,
         )
 
-        # Pixel values are used only in the first iteration if available
-        # In subsquent iterations, they are already merged with text and cached
-        # NOTE: first iteration doesn't have to be prefill, it can be the first
-        # iteration with a question and cached system prompt (continue generate from cache)
-        if is_first_iteration or not kwargs.get("use_cache", True):
+        # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
+        # Otherwise we need pixel values to be passed to model
+        if cache_position[0] == 0:
             model_inputs["pixel_values"] = pixel_values
 
         return model_inputs
@@ -1150,9 +1101,9 @@ class JanusForConditionalGeneration(JanusPreTrainedModel, GenerationMixin):
     @torch.no_grad()
     def generate(
         self,
-        inputs: torch.Tensor | None = None,
-        attention_mask: torch.LongTensor | None = None,
-        logits_processor: LogitsProcessorList | None = None,
+        inputs: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        logits_processor: Optional[LogitsProcessorList] = None,
         **kwargs,
     ):
         # 1. Handle generation config and model kwargs
@@ -1394,16 +1345,16 @@ class JanusImageProcessor(BlipImageProcessor):
     def __init__(
         self,
         do_resize: bool = True,
-        size: dict[str, int] | None = None,
+        size: Optional[dict[str, int]] = None,
         min_size: int = 14,
         resample: PILImageResampling = PILImageResampling.BICUBIC,
         do_rescale: bool = True,
-        rescale_factor: int | float = 1 / 255,
+        rescale_factor: Union[int, float] = 1 / 255,
         do_normalize: bool = True,
-        image_mean: float | list[float] | None = None,
-        image_std: float | list[float] | None = None,
-        do_convert_rgb: bool | None = None,
-        do_pad: bool | None = True,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        do_pad: Optional[bool] = True,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -1418,9 +1369,9 @@ class JanusImageProcessor(BlipImageProcessor):
     def pad_to_square(
         self,
         image: np.ndarray,
-        background_color: int | tuple[int, int, int] = 0,
-        data_format: str | ChannelDimension | None = None,
-        input_data_format: str | ChannelDimension | None = None,
+        background_color: Union[int, tuple[int, int, int]] = 0,
+        data_format: Optional[Union[str, ChannelDimension]] = None,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
     ) -> np.ndarray:
         """
         Pads an image to a square based on the longest edge.
@@ -1492,10 +1443,10 @@ class JanusImageProcessor(BlipImageProcessor):
     def resize(
         self,
         image: np.ndarray,
-        size: dict[str, int] | int,
+        size: Union[dict[str, int], int],
         resample: PILImageResampling = PILImageResampling.BICUBIC,
-        data_format: str | ChannelDimension | None = None,
-        input_data_format: str | ChannelDimension | None = None,
+        data_format: Optional[Union[str, ChannelDimension]] = None,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
         **kwargs,
     ) -> np.ndarray:
         """
@@ -1558,20 +1509,20 @@ class JanusImageProcessor(BlipImageProcessor):
     def preprocess(
         self,
         images: ImageInput,
-        do_resize: bool | None = None,
-        size: dict[str, int] | None = None,
-        resample: PILImageResampling | None = None,
-        do_rescale: bool | None = None,
-        rescale_factor: float | None = None,
-        do_normalize: bool | None = None,
-        image_mean: float | list[float] | None = None,
-        image_std: float | list[float] | None = None,
-        return_tensors: str | TensorType | None = None,
-        do_convert_rgb: bool | None = None,
-        background_color: int | tuple[int, int, int] | None = None,
-        do_pad: bool | None = None,
+        do_resize: Optional[bool] = None,
+        size: Optional[dict[str, int]] = None,
+        resample: Optional[PILImageResampling] = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        background_color: Optional[Union[int, tuple[int, int, int]]] = None,
+        do_pad: Optional[bool] = None,
         data_format: ChannelDimension = ChannelDimension.FIRST,
-        input_data_format: str | ChannelDimension | None = None,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
     ) -> PIL.Image.Image:
         """
         Preprocess an image or batch of images.
@@ -1708,13 +1659,13 @@ class JanusImageProcessor(BlipImageProcessor):
     def postprocess(
         self,
         images: ImageInput,
-        do_rescale: bool | None = None,
-        rescale_factor: float | None = None,
-        do_normalize: bool | None = None,
-        image_mean: list[float] | None = None,
-        image_std: list[float] | None = None,
-        input_data_format: str | None = None,
-        return_tensors: str | None = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[list[float]] = None,
+        image_std: Optional[list[float]] = None,
+        input_data_format: Optional[str] = None,
+        return_tensors: Optional[str] = None,
     ):
         """Applies post-processing to the decoded image tokens by reversing transformations applied during preprocessing."""
         do_rescale = do_rescale if do_rescale is not None else self.do_rescale
@@ -1759,9 +1710,9 @@ class JanusImageProcessor(BlipImageProcessor):
     def unnormalize(
         self,
         image: np.ndarray,
-        image_mean: float | Iterable[float],
-        image_std: float | Iterable[float],
-        input_data_format: str | ChannelDimension | None = None,
+        image_mean: Union[float, Iterable[float]],
+        image_std: Union[float, Iterable[float]],
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
     ) -> np.ndarray:
         """
         Unnormalizes `image` using the mean and standard deviation specified by `mean` and `std`.

transformers/models/janus/processing_janus.py

@@ -1,3 +1,4 @@
+# coding=utf-8
 # Copyright 2025 Deepseek AI and The HuggingFace Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,11 +16,13 @@
 Processor class for Janus.
 """
 
+from typing import Optional, Union
+
 from ...feature_extraction_utils import BatchFeature
 from ...image_utils import ImageInput
 from ...processing_utils import ProcessingKwargs, ProcessorMixin, TextKwargs, Unpack
 from ...tokenization_utils_base import PreTokenizedInput, TextInput
-from ...utils import auto_docstring, logging
+from ...utils import logging
 
 
 logger = logging.get_logger(__name__)
@@ -32,13 +35,6 @@ DEFAULT_SYSTEM_PROMPT = (
 
 
 class JanusTextKwargs(TextKwargs, total=False):
-    """
-    generation_mode (`str`, *optional*, defaults to `"text"`):
-        The generation mode indicating which modality to generate. Can be one of `"text"` or `"image"`. When set
-        to `"text"`, the processor prepares inputs for text generation. When set to `"image"`, it prepares inputs
-        for image generation by appending image start tokens to the prompt.
-    """
-
     generation_mode: str
 
 
@@ -50,13 +46,25 @@ class JanusProcessorKwargs(ProcessingKwargs, total=False):
     }
 
 
-@auto_docstring
 class JanusProcessor(ProcessorMixin):
-    def __init__(self, image_processor, tokenizer, chat_template=None, use_default_system_prompt=False, **kwargs):
-        r"""
-        use_default_system_prompt (`bool`, *optional*, defaults to `False`):
+    r"""
+    Constructs a Janus processor which wraps a Janus Image Processor and a Llama tokenizer into a single processor.
+
+    [`JanusProcessor`] offers all the functionalities of [`JanusImageProcessor`] and [`LlamaTokenizerFast`]. See the
+    [`~JanusProcessor.__call__`] and [`~JanusProcessor.decode`] for more information.
+
+    Args:
+        image_processor ([`JanusImageProcessor`]):
+            The image processor is a required input.
+        tokenizer ([`LlamaTokenizerFast`]):
+            The tokenizer is a required input.
+        chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
+            in a chat into a tokenizable string.
+        use_default_system_prompt (`str`, *optional*, defaults to `False`):
             Use default system prompt for Text Generation.
-        """
+    """
+
+    def __init__(self, image_processor, tokenizer, chat_template=None, use_default_system_prompt=False, **kwargs):
         self.num_image_tokens = 576
         self.image_token = tokenizer.image_token
         self.image_start_token = tokenizer.boi_token
@@ -65,14 +73,32 @@ class JanusProcessor(ProcessorMixin):
 
         super().__init__(image_processor, tokenizer, chat_template=chat_template)
 
-    @auto_docstring
     def __call__(
         self,
-        text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None,
-        images: ImageInput | None = None,
+        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,
+        images: Optional[ImageInput] = None,
         **kwargs: Unpack[JanusProcessorKwargs],
     ) -> BatchFeature:
-        r"""
+        """
+        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+        and `kwargs` arguments to LlamaTokenizerFast's [`~LlamaTokenizerFast.__call__`] if `text` is not `None` to encode
+        the text. To prepare the image(s), this method forwards the `images` and `kwargs` arguments to
+        JanusImageProcessor's [`~JanusImageProcessor.__call__`] if `images` is not `None`. Please refer to the doctsring
+        of the above two methods for more information.
+
+        Args:
+            text (`str`, `list[str]`, `list[list[str]]`):
+                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):
+                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+                tensor. Both channels-first and channels-last formats are supported.
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors of a particular framework. Acceptable values are:
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return NumPy `np.ndarray` objects.
+
         Returns:
             [`BatchFeature`]: A [`BatchFeature`] with the following fields:
 

transformers/models/jetmoe/configuration_jetmoe.py

@@ -1,3 +1,4 @@
+# coding=utf-8
 # Copyright 2024 JetMoe AI and the HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,6 +14,8 @@
 # limitations under the License.
 """JetMoe model configuration"""
 
+from typing import Optional
+
 from ...configuration_utils import PreTrainedConfig
 from ...modeling_rope_utils import RopeParameters
 from ...utils import logging
@@ -100,26 +103,26 @@ class JetMoeConfig(PreTrainedConfig):
 
     def __init__(
         self,
-        vocab_size: int | None = 32000,
-        hidden_size: int | None = 2048,
-        num_hidden_layers: int | None = 12,
-        num_key_value_heads: int | None = 16,
-        kv_channels: int | None = 128,
-        intermediate_size: int | None = 5632,
-        max_position_embeddings: int | None = 4096,
-        activation_function: str | None = "silu",
-        num_local_experts: int | None = 8,
-        num_experts_per_tok: int | None = 2,
-        output_router_logits: bool | None = False,
-        aux_loss_coef: float | None = 0.01,
-        use_cache: bool | None = True,
-        bos_token_id: int | None = 1,
-        eos_token_id: int | None = 2,
-        tie_word_embeddings: bool | None = True,
-        rope_parameters: RopeParameters | dict[str, RopeParameters] | None = None,
-        rms_norm_eps: int | None = 1e-6,
-        initializer_range: float | None = 0.01,
-        attention_dropout: float | None = 0.0,
+        vocab_size: Optional[int] = 32000,
+        hidden_size: Optional[int] = 2048,
+        num_hidden_layers: Optional[int] = 12,
+        num_key_value_heads: Optional[int] = 16,
+        kv_channels: Optional[int] = 128,
+        intermediate_size: Optional[int] = 5632,
+        max_position_embeddings: Optional[int] = 4096,
+        activation_function: Optional[str] = "silu",
+        num_local_experts: Optional[int] = 8,
+        num_experts_per_tok: Optional[int] = 2,
+        output_router_logits: Optional[bool] = False,
+        aux_loss_coef: Optional[float] = 0.01,
+        use_cache: Optional[bool] = True,
+        bos_token_id: Optional[int] = 1,
+        eos_token_id: Optional[int] = 2,
+        tie_word_embeddings: Optional[bool] = True,
+        rope_parameters: Optional[RopeParameters | dict[str, RopeParameters]] = None,
+        rms_norm_eps: Optional[int] = 1e-6,
+        initializer_range: Optional[float] = 0.01,
+        attention_dropout: Optional[float] = 0.0,
         **kwargs,
     ):
         if num_experts_per_tok > num_local_experts:
@@ -146,10 +149,9 @@ class JetMoeConfig(PreTrainedConfig):
         self.rms_norm_eps = rms_norm_eps
         self.rope_parameters = rope_parameters
 
-        self.bos_token_id = bos_token_id
-        self.eos_token_id = eos_token_id
-        self.tie_word_embeddings = tie_word_embeddings
-        super().__init__(**kwargs)
+        super().__init__(
+            bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs
+        )
 
 
 __all__ = ["JetMoeConfig"]