optimum-rbln 0.8.0.post2__py3-none-any.whl → 0.8.1a1__py3-none-any.whl

optimum/rbln/transformers/configuration_generic.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from ..configuration_utils import RBLNModelConfig
 
 
-class _RBLNTransformerEncoderConfig(RBLNModelConfig):
+class RBLNTransformerEncoderConfig(RBLNModelConfig):
     rbln_model_input_names: Optional[List[str]] = None
 
     def __init__(
@@ -25,7 +25,7 @@ class _RBLNTransformerEncoderConfig(RBLNModelConfig):
         max_seq_len: Optional[int] = None,
         batch_size: Optional[int] = None,
         model_input_names: Optional[List[str]] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -47,9 +47,12 @@ class _RBLNTransformerEncoderConfig(RBLNModelConfig):
         self.model_input_names = model_input_names or self.rbln_model_input_names
 
 
-class _RBLNImageModelConfig(RBLNModelConfig):
+class RBLNImageModelConfig(RBLNModelConfig):
     def __init__(
-        self, image_size: Optional[Union[int, Tuple[int, int]]] = None, batch_size: Optional[int] = None, **kwargs
+        self,
+        image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        batch_size: Optional[int] = None,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -86,32 +89,32 @@ class _RBLNImageModelConfig(RBLNModelConfig):
             return self.image_size["height"]
 
 
-class RBLNModelForQuestionAnsweringConfig(_RBLNTransformerEncoderConfig):
+class RBLNModelForQuestionAnsweringConfig(RBLNTransformerEncoderConfig):
     pass
 
 
-class RBLNModelForSequenceClassificationConfig(_RBLNTransformerEncoderConfig):
+class RBLNModelForSequenceClassificationConfig(RBLNTransformerEncoderConfig):
     pass
 
 
-class RBLNModelForMaskedLMConfig(_RBLNTransformerEncoderConfig):
+class RBLNModelForMaskedLMConfig(RBLNTransformerEncoderConfig):
     pass
 
 
-class RBLNModelForTextEncodingConfig(_RBLNTransformerEncoderConfig):
+class RBLNModelForTextEncodingConfig(RBLNTransformerEncoderConfig):
     pass
 
 
 # FIXME : Appropriate name ?
-class RBLNTransformerEncoderForFeatureExtractionConfig(_RBLNTransformerEncoderConfig):
+class RBLNTransformerEncoderForFeatureExtractionConfig(RBLNTransformerEncoderConfig):
     pass
 
 
-class RBLNModelForImageClassificationConfig(_RBLNImageModelConfig):
+class RBLNModelForImageClassificationConfig(RBLNImageModelConfig):
     pass
 
 
-class RBLNModelForDepthEstimationConfig(_RBLNImageModelConfig):
+class RBLNModelForDepthEstimationConfig(RBLNImageModelConfig):
     pass
 
 
@@ -121,7 +124,7 @@ class RBLNModelForAudioClassificationConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         max_length: Optional[int] = None,
         num_mel_bins: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

optimum/rbln/transformers/modeling_generic.py

@@ -43,9 +43,9 @@ from ..configuration_utils import RBLNCompileConfig
 from ..modeling import RBLNModel
 from ..utils.logging import get_logger
 from .configuration_generic import (
+    RBLNImageModelConfig,
     RBLNModelForAudioClassificationConfig,
-    _RBLNImageModelConfig,
-    _RBLNTransformerEncoderConfig,
+    RBLNTransformerEncoderConfig,
 )
 
 
@@ -55,7 +55,7 @@ if TYPE_CHECKING:
 logger = get_logger()
 
 
-class _RBLNTransformerEncoder(RBLNModel):
+class RBLNTransformerEncoder(RBLNModel):
     auto_model_class = AutoModel
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
     rbln_dtype = "int64"
@@ -66,8 +66,8 @@ class _RBLNTransformerEncoder(RBLNModel):
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
         model: Optional["PreTrainedModel"] = None,
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_config: Optional[_RBLNTransformerEncoderConfig] = None,
-    ) -> _RBLNTransformerEncoderConfig:
+        rbln_config: Optional[RBLNTransformerEncoderConfig] = None,
+    ) -> RBLNTransformerEncoderConfig:
         return cls.update_rbln_config_for_transformers_encoder(
             preprocessors=preprocessors,
             model=model,
@@ -81,8 +81,8 @@ class _RBLNTransformerEncoder(RBLNModel):
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
         model: Optional["PreTrainedModel"] = None,
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_config: Optional[_RBLNTransformerEncoderConfig] = None,
-    ) -> _RBLNTransformerEncoderConfig:
+        rbln_config: Optional[RBLNTransformerEncoderConfig] = None,
+    ) -> RBLNTransformerEncoderConfig:
         max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
             model_config, "max_position_embeddings", None
         )
@@ -150,8 +150,8 @@ class _RBLNImageModel(RBLNModel):
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
         model: Optional["PreTrainedModel"] = None,
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_config: Optional[_RBLNImageModelConfig] = None,
-    ) -> _RBLNImageModelConfig:
+        rbln_config: Optional[RBLNImageModelConfig] = None,
+    ) -> RBLNImageModelConfig:
         return cls.update_rbln_config_for_image_model(
             preprocessors=preprocessors,
             model=model,
@@ -165,8 +165,8 @@ class _RBLNImageModel(RBLNModel):
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
         model: Optional["PreTrainedModel"] = None,
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_config: Optional[_RBLNImageModelConfig] = None,
-    ) -> _RBLNImageModelConfig:
+        rbln_config: Optional[RBLNImageModelConfig] = None,
+    ) -> RBLNImageModelConfig:
         if rbln_config.image_size is None:
             for processor in preprocessors:
                 if hasattr(processor, "size"):
@@ -196,15 +196,14 @@ class _RBLNImageModel(RBLNModel):
         return rbln_config
 
 
-class RBLNModelForQuestionAnswering(_RBLNTransformerEncoder):
+class RBLNModelForQuestionAnswering(RBLNTransformerEncoder):
     auto_model_class = AutoModelForQuestionAnswering
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
     output_class = QuestionAnsweringModelOutput
 
     def _prepare_output(self, output, return_dict):
-        """
-        Prepare QuestionAnswering specific output format.
-        """
+        # Prepare QuestionAnswering specific output format.
+
         start_logits, end_logits = output
 
         if not return_dict:
@@ -213,22 +212,22 @@ class RBLNModelForQuestionAnswering(_RBLNTransformerEncoder):
             return QuestionAnsweringModelOutput(start_logits=start_logits, end_logits=end_logits)
 
 
-class RBLNModelForSequenceClassification(_RBLNTransformerEncoder):
+class RBLNModelForSequenceClassification(RBLNTransformerEncoder):
     auto_model_class = AutoModelForSequenceClassification
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
-class RBLNModelForMaskedLM(_RBLNTransformerEncoder):
+class RBLNModelForMaskedLM(RBLNTransformerEncoder):
     auto_model_class = AutoModelForMaskedLM
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
-class RBLNModelForTextEncoding(_RBLNTransformerEncoder):
+class RBLNModelForTextEncoding(RBLNTransformerEncoder):
     auto_model_class = AutoModelForTextEncoding
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
-class RBLNTransformerEncoderForFeatureExtraction(_RBLNTransformerEncoder):
+class RBLNTransformerEncoderForFeatureExtraction(RBLNTransformerEncoder):
     # TODO: RBLNModel is also for feature extraction.
     auto_model_class = AutoModel
     rbln_model_input_names = ["input_ids", "attention_mask"]

optimum/rbln/transformers/modeling_rope_utils.py

@@ -51,7 +51,7 @@ def _compute_default_rope_parameters(
 
     base = config.rope_theta
     partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
-    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
     dim = int(head_dim * partial_rotary_factor)
 
     attention_factor = 1.0  # Unused in this type of RoPE

optimum/rbln/transformers/models/__init__.py

@@ -18,6 +18,10 @@ from transformers.utils import _LazyModule
 
 
 _import_structure = {
+    "audio_spectrogram_transformer": [
+        "RBLNASTForAudioClassification",
+        "RBLNASTForAudioClassificationConfig",
+    ],
     "auto": [
         "RBLNAutoModel",
         "RBLNAutoModelForAudioClassification",
@@ -65,6 +69,10 @@ _import_structure = {
         "RBLNCLIPVisionModelWithProjection",
         "RBLNCLIPVisionModelWithProjectionConfig",
     ],
+    "distilbert": [
+        "RBLNDistilBertForQuestionAnswering",
+        "RBLNDistilBertForQuestionAnsweringConfig",
+    ],
     "qwen2_5_vl": [
         "RBLNQwen2_5_VisionTransformerPretrainedModel",
         "RBLNQwen2_5_VisionTransformerPretrainedModelConfig",
@@ -101,11 +109,18 @@ _import_structure = {
     "mistral": ["RBLNMistralForCausalLM", "RBLNMistralForCausalLMConfig"],
     "phi": ["RBLNPhiForCausalLM", "RBLNPhiForCausalLMConfig"],
     "qwen2": ["RBLNQwen2ForCausalLM", "RBLNQwen2ForCausalLMConfig"],
+    "resnet": ["RBLNResNetForImageClassification", "RBLNResNetForImageClassificationConfig"],
+    "roberta": [
+        "RBLNRobertaForMaskedLM",
+        "RBLNRobertaForMaskedLMConfig",
+        "RBLNRobertaForSequenceClassification",
+        "RBLNRobertaForSequenceClassificationConfig",
+    ],
     "siglip": [
         "RBLNSiglipVisionModel",
         "RBLNSiglipVisionModelConfig",
     ],
-    "time_series_transformers": [
+    "time_series_transformer": [
         "RBLNTimeSeriesTransformerForPrediction",
         "RBLNTimeSeriesTransformerForPredictionConfig",
     ],
@@ -115,12 +130,22 @@ _import_structure = {
         "RBLNT5EncoderModelConfig",
         "RBLNT5ForConditionalGenerationConfig",
     ],
+    "vit": ["RBLNViTForImageClassification", "RBLNViTForImageClassificationConfig"],
     "wav2vec2": ["RBLNWav2Vec2ForCTC", "RBLNWav2Vec2ForCTCConfig"],
     "whisper": ["RBLNWhisperForConditionalGeneration", "RBLNWhisperForConditionalGenerationConfig"],
-    "xlm_roberta": ["RBLNXLMRobertaModel", "RBLNXLMRobertaModelConfig"],
+    "xlm_roberta": [
+        "RBLNXLMRobertaModel",
+        "RBLNXLMRobertaModelConfig",
+        "RBLNXLMRobertaForSequenceClassification",
+        "RBLNXLMRobertaForSequenceClassificationConfig",
+    ],
 }
 
 if TYPE_CHECKING:
+    from .audio_spectrogram_transformer import (
+        RBLNASTForAudioClassification,
+        RBLNASTForAudioClassificationConfig,
+    )
     from .auto import (
         RBLNAutoModel,
         RBLNAutoModelForAudioClassification,
@@ -172,6 +197,10 @@ if TYPE_CHECKING:
         RBLNDecoderOnlyModelForCausalLM,
         RBLNDecoderOnlyModelForCausalLMConfig,
     )
+    from .distilbert import (
+        RBLNDistilBertForQuestionAnswering,
+        RBLNDistilBertForQuestionAnsweringConfig,
+    )
     from .dpt import (
         RBLNDPTForDepthEstimation,
         RBLNDPTForDepthEstimationConfig,
@@ -204,6 +233,13 @@ if TYPE_CHECKING:
         RBLNQwen2_5_VLForConditionalGeneration,
         RBLNQwen2_5_VLForConditionalGenerationConfig,
     )
+    from .resnet import RBLNResNetForImageClassification, RBLNResNetForImageClassificationConfig
+    from .roberta import (
+        RBLNRobertaForMaskedLM,
+        RBLNRobertaForMaskedLMConfig,
+        RBLNRobertaForSequenceClassification,
+        RBLNRobertaForSequenceClassificationConfig,
+    )
     from .siglip import RBLNSiglipVisionModel, RBLNSiglipVisionModelConfig
     from .t5 import (
         RBLNT5EncoderModel,
@@ -211,13 +247,19 @@ if TYPE_CHECKING:
         RBLNT5ForConditionalGeneration,
         RBLNT5ForConditionalGenerationConfig,
     )
-    from .time_series_transformers import (
+    from .time_series_transformer import (
         RBLNTimeSeriesTransformerForPrediction,
         RBLNTimeSeriesTransformerForPredictionConfig,
     )
+    from .vit import RBLNViTForImageClassification, RBLNViTForImageClassificationConfig
     from .wav2vec2 import RBLNWav2Vec2ForCTC, RBLNWav2Vec2ForCTCConfig
     from .whisper import RBLNWhisperForConditionalGeneration, RBLNWhisperForConditionalGenerationConfig
-    from .xlm_roberta import RBLNXLMRobertaModel, RBLNXLMRobertaModelConfig
+    from .xlm_roberta import (
+        RBLNXLMRobertaForSequenceClassification,
+        RBLNXLMRobertaForSequenceClassificationConfig,
+        RBLNXLMRobertaModel,
+        RBLNXLMRobertaModelConfig,
+    )
 
 else:
     import sys

optimum/rbln/transformers/models/audio_spectrogram_transformer/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from .configuration_audio_spectrogram_transformer import RBLNASTForAudioClassificationConfig
+from .modeling_audio_spectrogram_transformer import RBLNASTForAudioClassification

optimum/rbln/transformers/models/audio_spectrogram_transformer/configuration_audio_spectrogram_transformer.py

@@ -0,0 +1,21 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...configuration_generic import RBLNModelForAudioClassificationConfig
+
+
+class RBLNASTForAudioClassificationConfig(RBLNModelForAudioClassificationConfig):
+    """
+    Configuration class for RBLNASTForAudioClassification.
+    """

optimum/rbln/transformers/models/audio_spectrogram_transformer/modeling_audio_spectrogram_transformer.py

@@ -0,0 +1,28 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...modeling_generic import RBLNModelForAudioClassification
+
+
+class RBLNASTForAudioClassification(RBLNModelForAudioClassification):
+    """
+    Audio Spectrogram Transformer model with an audio classification head on top (a linear layer on top of the pooled output) e.g. for datasets like AudioSet, Speech Commands v2.
+    This model inherits from [`RBLNModelForAudioClassification`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformer-based `ASTForAudioClassification` models on RBLN devices.
+    It implements the methods to convert a pre-trained transformers `ASTForAudioClassification` model into a RBLN transformer model by:
+
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN Compiler.
+    """

optimum/rbln/transformers/models/auto/auto_factory.py

@@ -11,10 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 import importlib
 import inspect
 import warnings
+from typing import Type
 
 from transformers import AutoConfig, PretrainedConfig
 from transformers.dynamic_module_utils import get_class_from_dynamic_module
@@ -22,7 +22,12 @@ from transformers.models.auto.auto_factory import _get_model_class
 
 from optimum.rbln.configuration_utils import RBLNAutoConfig
 from optimum.rbln.modeling_base import RBLNBaseModel
-from optimum.rbln.utils.model_utils import convert_hf_to_rbln_model_name, convert_rbln_to_hf_model_name
+from optimum.rbln.utils.model_utils import (
+    MODEL_MAPPING,
+    convert_hf_to_rbln_model_name,
+    convert_rbln_to_hf_model_name,
+    get_rbln_model_cls,
+)
 
 
 class _BaseAutoModelClass:
@@ -58,7 +63,7 @@ class _BaseAutoModelClass:
             hf_model_class = cls.infer_hf_model_class(pretrained_model_name_or_path, **kwargs)
             rbln_class_name = convert_hf_to_rbln_model_name(hf_model_class.__name__)
         else:
-            rbln_class_name = cls.get_rbln_model_class_name(pretrained_model_name_or_path, **kwargs)
+            rbln_class_name = cls.get_rbln_model_cls_name(pretrained_model_name_or_path, **kwargs)
 
             if convert_rbln_to_hf_model_name(rbln_class_name) not in cls._model_mapping_names.values():
                 raise ValueError(
@@ -68,8 +73,7 @@ class _BaseAutoModelClass:
                 )
 
         try:
-            module = importlib.import_module("optimum.rbln")
-            rbln_cls = getattr(module, rbln_class_name)
+            rbln_cls = get_rbln_model_cls(rbln_class_name)
         except AttributeError as e:
             raise AttributeError(
                 f"Class '{rbln_class_name}' not found in 'optimum.rbln' module for model ID '{pretrained_model_name_or_path}'. "
@@ -136,7 +140,7 @@ class _BaseAutoModelClass:
         return model_class
 
     @classmethod
-    def get_rbln_model_class_name(cls, pretrained_model_name_or_path, **kwargs):
+    def get_rbln_model_cls_name(cls, pretrained_model_name_or_path, **kwargs):
         """
         Retrieve the path to the compiled model directory for a given RBLN model.
 
@@ -159,11 +163,25 @@ class _BaseAutoModelClass:
         return rbln_config.rbln_model_cls_name
 
     @classmethod
-    def from_pretrained(
-        cls,
-        model_id,
-        *args,
-        **kwargs,
-    ):
+    def from_pretrained(cls, model_id, *args, **kwargs):
         rbln_cls = cls.get_rbln_cls(model_id, *args, **kwargs)
         return rbln_cls.from_pretrained(model_id, *args, **kwargs)
+
+    @staticmethod
+    def register(rbln_cls: Type[RBLNBaseModel], exist_ok=False):
+        """
+        Register a new RBLN model class.
+
+        Args:
+            rbln_cls (Type[RBLNBaseModel]): The RBLN model class to register.
+            exist_ok (bool): Whether to allow registering an already registered model.
+        """
+        if not issubclass(rbln_cls, RBLNBaseModel):
+            raise ValueError("`rbln_cls` must be a subclass of RBLNBaseModel.")
+
+        native_cls = getattr(importlib.import_module("optimum.rbln"), rbln_cls.__name__, None)
+        if rbln_cls.__name__ in MODEL_MAPPING or native_cls is not None:
+            if not exist_ok:
+                raise ValueError(f"Model for {rbln_cls.__name__} already registered.")
+
+        MODEL_MAPPING[rbln_cls.__name__] = rbln_cls

optimum/rbln/transformers/models/blip_2/modeling_blip_2.py

@@ -251,6 +251,38 @@ class RBLNBlip2QFormerModel(RBLNModel):
 
 
 class RBLNBlip2ForConditionalGeneration(RBLNModel):
+    """
+    RBLNBlip2ForConditionalGeneration is a multi-modal model that integrates vision and language processing capabilities,
+    optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.
+
+    This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    Important Note:
+        This model includes a Large Language Model (LLM) as a submodule. For optimal performance, it is highly recommended to use
+        tensor parallelism for the language model.  This can be achieved by using the `rbln_config` parameter in the
+        `from_pretrained` method. Refer to the `from_pretrained` documentation and the RBLNBlip2ForConditionalGeneration class for details.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNBlip2ForConditionalGeneration
+
+        model = RBLNBlip2ForConditionalGeneration.from_pretrained(
+            "Salesforce/blip2-opt-2.7b",
+            export=True,
+            rbln_config={
+                "language_model": {
+                    "batch_size": 1,
+                    "max_seq_len": 2048,
+                    "tensor_parallel_size": 1,
+                    "use_inputs_embeds": True,
+                },
+            },
+        )
+
+        model.save_pretrained("compiled-blip2-opt-2.7b")
+        ```
+    """
+
     auto_model_class = AutoModelForVisualQuestionAnswering
     _rbln_submodules = [{"name": "vision_model"}, {"name": "qformer"}, {"name": "language_model"}]
 
@@ -275,10 +307,9 @@ class RBLNBlip2ForConditionalGeneration(RBLNModel):
         subfolder: str,
         rbln_config: RBLNModelConfig,
     ):
-        """
-        If you are unavoidably running on a CPU rather than an RBLN device,
-        store the torch tensor, weight, etc. in this function.
-        """
+        # If you are unavoidably running on a CPU rather than an RBLN device,
+        # store the torch tensor, weight, etc. in this function.
+
         save_dict = {}
         save_dict["query_tokens"] = model.query_tokens
         torch.save(save_dict, save_dir_path / subfolder / "query_tokens.pth")

optimum/rbln/transformers/models/clip/configuration_clip.py

@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNCLIPTextModelConfig(RBLNModelConfig):
-    def __init__(self, batch_size: Optional[int] = None, **kwargs):
+    def __init__(self, batch_size: Optional[int] = None, **kwargs: Dict[str, Any]):
         """
         Args:
             batch_size (Optional[int]): The batch size for text processing. Defaults to 1.
@@ -38,7 +38,7 @@ class RBLNCLIPTextModelWithProjectionConfig(RBLNCLIPTextModelConfig):
 
 
 class RBLNCLIPVisionModelConfig(RBLNModelConfig):
-    def __init__(self, batch_size: Optional[int] = None, image_size: Optional[int] = None, **kwargs):
+    def __init__(self, batch_size: Optional[int] = None, image_size: Optional[int] = None, **kwargs: Dict[str, Any]):
         """
         Args:
             batch_size (Optional[int]): The batch size for image processing. Defaults to 1.

optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -49,7 +49,7 @@ class RBLNCLIPTextModel(RBLNModel):
 
     @classmethod
     def update_rbln_config_using_pipe(
-        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_config: str
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
     ) -> "RBLNDiffusionMixinConfig":
         return rbln_config
 
@@ -81,10 +81,9 @@ class RBLNCLIPTextModel(RBLNModel):
         return output
 
     def _prepare_output(self, output, return_dict):
-        """
-        Prepare model output based on return_dict flag.
-        This method can be overridden by subclasses to provide task-specific output handling.
-        """
+        # Prepare model output based on return_dict flag.
+        # This method can be overridden by subclasses to provide task-specific output handling.
+
         if not return_dict:
             return (output,) if not isinstance(output, (tuple, list)) else output
         else:
@@ -161,17 +160,17 @@ class RBLNCLIPVisionModel(RBLNModel):
         return_dict: bool = None,
         **kwargs,
     ) -> Union[Tuple, CLIPVisionModelOutput]:
-        if len(kwargs) > 0 and any(kwargs.values()):
-            logger.warning(f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__}.")
-
+        if len(kwargs) > 0 and any(value is not None for value in kwargs.values()):
+            logger.warning(
+                f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__.__name__}."
+            )
         output = super().forward(pixel_values, return_dict=return_dict)
         return output
 
     def _prepare_output(self, output, return_dict):
-        """
-        Prepare model output based on return_dict flag.
-        This method can be overridden by subclasses to provide task-specific output handling.
-        """
+        # Prepare model output based on return_dict flag.
+        # This method can be overridden by subclasses to provide task-specific output handling.
+
         if not return_dict:
             return (output,) if not isinstance(output, (tuple, list)) else output
         else: