optimum-rbln 0.8.1a0__py3-none-any.whl → 0.8.1a2__py3-none-any.whl

optimum/rbln/transformers/models/siglip/modeling_siglip.py

@@ -12,12 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 import torch
 from transformers import SiglipVisionConfig, SiglipVisionModel
 from transformers.modeling_outputs import BaseModelOutputWithPooling
-from transformers.models.siglip.modeling_siglip import SiglipVisionModelOutput
 
 from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
@@ -34,11 +33,18 @@ if TYPE_CHECKING:
 
 
 class _SiglipVisionModel(torch.nn.Module):
-    def __init__(self, model: SiglipVisionModel, interpolate_pos_encoding: bool, output_hidden_states: bool):
+    def __init__(
+        self,
+        model: SiglipVisionModel,
+        interpolate_pos_encoding: bool,
+        output_hidden_states: bool,
+        output_attentions: bool,
+    ):
         super().__init__()
         self.vision_model = model.vision_model
         self.interpolate_pos_encoding = interpolate_pos_encoding
         self.output_hidden_states = output_hidden_states
+        self.output_attentions = output_attentions
 
     def forward(self, inp):
         enc_out = self.vision_model(
@@ -46,6 +52,7 @@ class _SiglipVisionModel(torch.nn.Module):
             output_hidden_states=self.output_hidden_states,
             return_dict=False,
             interpolate_pos_encoding=self.interpolate_pos_encoding,
+            output_attentions=self.output_attentions,
         )
         return tuple(x for x in enc_out if x is not None)
 
@@ -56,6 +63,7 @@ class RBLNSiglipVisionModel(RBLNModel):
         wrapper_cfg = {
             "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
             "output_hidden_states": rbln_config.output_hidden_states,
+            "output_attentions": rbln_config.output_attentions,
         }
         return _SiglipVisionModel(model, **wrapper_cfg).eval()
 
@@ -81,8 +89,10 @@ class RBLNSiglipVisionModel(RBLNModel):
         if rbln_config.image_size is None:
             raise ValueError("`rbln_image_size` should be specified!")
 
+        if rbln_config.output_attentions is None:
+            rbln_config.output_attentions = getattr(model_config, "output_attentions", False)
         if rbln_config.output_hidden_states is None:
-            rbln_config.output_hidden_states = model_config.output_hidden_states
+            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
 
         rbln_compile_config = RBLNCompileConfig(
             input_info=[
@@ -104,43 +114,74 @@ class RBLNSiglipVisionModel(RBLNModel):
 
     def forward(
         self,
-        pixel_values: Optional[torch.FloatTensor] = None,
+        pixel_values: torch.Tensor,
         return_dict: bool = None,
+        output_attentions: bool = None,
+        output_hidden_states: bool = None,
         interpolate_pos_encoding: bool = False,
-        **kwargs,
-    ) -> Union[Tuple, SiglipVisionModelOutput]:
-        if len(kwargs) > 0 and any(kwargs.values()):
-            logger.warning(f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__}.")
+        **kwargs: Dict[str, Any],
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        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_attentions = output_attentions if output_attentions is not None else self.rbln_config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        if output_attentions != self.rbln_config.output_attentions:
+            raise ValueError(
+                f"Variable output_attentions {output_attentions} is not equal to rbln_config.output_attentions {self.rbln_config.output_attentions} "
+                f"Please compile again with the correct argument."
+            )
+
+        if output_hidden_states != self.rbln_config.output_hidden_states:
+            raise ValueError(
+                f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {self.rbln_config.output_hidden_states} "
+                f"Please compile again with the correct argument."
+            )
 
         if interpolate_pos_encoding != self.rbln_config.interpolate_pos_encoding:
             raise ValueError(
-                f"Variable interpolate_pos_encoding {interpolate_pos_encoding} is not equal to rbln_config.interpolate_pos_encoding {self.rbln_config.interpolate_pos_encoding}"
+                f"Variable interpolate_pos_encoding {interpolate_pos_encoding} is not equal to rbln_config.interpolate_pos_encoding {self.rbln_config.interpolate_pos_encoding} "
                 f"Please compile again with the correct argument."
             )
+
         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:
-            last_hidden_state = (
-                output[0]
-                if self.rbln_config.interpolate_pos_encoding or self.rbln_config.output_hidden_states
-                else output
-            )
-            pooler_output = output[1] if self.rbln_config.interpolate_pos_encoding else None
+            last_hidden_state = output.pop(0) if isinstance(output, (tuple, list)) else output
+            vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+            pooler_output = output.pop(0) if getattr(vision_config, "vision_use_head", True) else None
+
             if self.rbln_config.output_hidden_states:
-                hidden_states = (output[2:] if self.rbln_config.interpolate_pos_encoding else output[1:],)
+                hidden_states = ()
+                num_hidden_layers = vision_config.num_hidden_layers
+                for _ in range(num_hidden_layers + 1):
+                    hidden_states += (output.pop(0),)
             else:
                 hidden_states = None
 
+            if self.rbln_config.output_attentions:
+                attentions = ()
+                num_hidden_layers = vision_config.num_hidden_layers
+                for _ in range(num_hidden_layers):
+                    attentions += (output.pop(0),)
+            else:
+                attentions = None
+
             return BaseModelOutputWithPooling(
                 last_hidden_state=last_hidden_state,
                 pooler_output=pooler_output,
                 hidden_states=hidden_states,
+                attentions=attentions,
             )

optimum/rbln/transformers/models/t5/modeling_t5.py

@@ -41,6 +41,28 @@ class T5EncoderWrapper(torch.nn.Module):
 
 
 class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
+    """
+    The T5 Model transformer with an encoder-only architecture for feature extraction.
+    This model inherits from [`RBLNTransformerEncoderForFeatureExtraction`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    Important Note:
+        This model supports various sizes of the T5EncoderModel. For optimal performance, it is highly recommended to adjust the tensor parallelism setting
+        based on the model size. Please refer to the [Optimum RBLN Overview](../../../optimum_rbln.md) for guidance on choosing the appropriate tensor parallelism size for your model.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNT5EncoderModel
+
+        model = RBLNT5EncoderModel.from_pretrained(
+            "sentence-transformers/sentence-t5-xxl",
+            export=True,
+            rbln_tensor_parallel_size=4,
+        )
+
+        model.save_pretrained("compiled-sentence-t5-xxl")
+        ```
+    """
+
     auto_model_class = AutoModelForTextEncoding
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
@@ -50,10 +72,7 @@ class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
 
     @classmethod
     def update_rbln_config_using_pipe(
-        cls,
-        pipe: "RBLNDiffusionMixin",
-        rbln_config: "RBLNDiffusionMixinConfig",
-        submodule_name: str,
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
     ) -> "RBLNDiffusionMixinConfig":
         submodule_config = getattr(rbln_config, submodule_name)
         submodule_config.max_seq_len = rbln_config.max_seq_len or 256
@@ -62,6 +81,29 @@ class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
 
 
 class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
+    """
+    The T5 Model transformer with a language modeling head for conditional generation.
+    This model inherits from [`RBLNModelForSeq2SeqLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    Important Note:
+        This model supports various sizes of the T5ForConditionalGeneration. For optimal performance, it is highly recommended to adjust the tensor parallelism setting
+        based on the model size. Please refer to the [Optimum RBLN Overview](../../../optimum_rbln.md) for guidance on choosing the appropriate tensor parallelism size for your model.
+
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNT5ForConditionalGeneration
+
+        model = RBLNT5ForConditionalGeneration.from_pretrained(
+            "google-t5/t5-11b",
+            export=True,
+            rbln_tensor_parallel_size=4,
+        )
+
+        model.save_pretrained("compiled-sentence-t5-xxl")
+        ```
+    """
+
     support_causal_attn = False
 
     @classmethod

optimum/rbln/transformers/models/t5/t5_architecture.py

@@ -136,10 +136,14 @@ class T5Decoder(Seq2SeqDecoder):
 
 
 class T5Block(Seq2SeqDecoderLayer):
+    def __init__(self, decoder_layer, self_attn):
+        super().__init__(decoder_layer, self_attn, cross_attn=None)
+        self.__post_init__()
+
     def __post_init__(self):
         self.self_attn_layer_norm = self._original_mod.layer[0].layer_norm
         self.encoder_attn_layer_norm = self._original_mod.layer[1].layer_norm
-        self.encoder_attn = T5CrossAttention(self._original_mod.layer[1].EncDecAttention)
+        self.cross_attn = T5CrossAttention(self._original_mod.layer[1].EncDecAttention)
         self.ff_layer = self._original_mod.layer[2]
 
     def pre_self_attn_layer_norm(self, hidden_states):

optimum/rbln/transformers/models/{time_series_transformers → time_series_transformer}/__init__.py

@@ -23,4 +23,4 @@
 
 from ....ops import paged_add_softmax_attn_decode, rbln_cache_update
 from .configuration_time_series_transformer import RBLNTimeSeriesTransformerForPredictionConfig
-from .modeling_time_series_transformers import RBLNTimeSeriesTransformerForPrediction
+from .modeling_time_series_transformer import RBLNTimeSeriesTransformerForPrediction

optimum/rbln/transformers/models/{time_series_transformers → time_series_transformer}/configuration_time_series_transformer.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -10,7 +10,7 @@ class RBLNTimeSeriesTransformerForPredictionConfig(RBLNModelConfig):
         enc_max_seq_len: Optional[int] = None,
         dec_max_seq_len: Optional[int] = None,
         num_parallel_samples: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

optimum/rbln/transformers/models/{time_series_transformers/modeling_time_series_transformers.py → time_series_transformer/modeling_time_series_transformer.py}

@@ -120,6 +120,17 @@ class RBLNSeq2SeqTSDecoderOutput(ModelOutput):
 
 
 class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
+    """
+    The Time Series Transformer Model with a distribution head on top for time-series forecasting. e.g., for datasets like M4, NN5, or other time series forecasting benchmarks.
+    This model inherits from [`RBLNModel`]. 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 `TimeSeriesTransformerForPrediction` models on RBLN devices.
+    It implements the methods to convert a pre-trained transformers `TimeSeriesTransformerForPrediction` 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.
+    """
+
     auto_model_class = None
     main_input_name = "inputs_embeds"
 
@@ -144,11 +155,6 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
         )
 
     def __getattr__(self, __name: str) -> Any:
-        """This is the key method to implement RBLN-TimeSeriesTransformersForPrediction.
-        Returns:
-            Any: TimeSeriesTransformersForPrediction's corresponding method
-        """
-
         def redirect(func):
             return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
 
@@ -211,10 +217,9 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
         subfolder: str,
         rbln_config: RBLNTimeSeriesTransformerForPredictionConfig,
     ):
-        """
-        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["embedder"] = model.model.embedder.state_dict()
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")

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

@@ -0,0 +1,19 @@
+# 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_vit import RBLNViTForImageClassificationConfig
+from .modeling_vit import RBLNViTForImageClassification
+
+
+__all__ = ["RBLNViTForImageClassificationConfig", "RBLNViTForImageClassification"]

optimum/rbln/transformers/models/vit/configuration_vit.py

@@ -0,0 +1,19 @@
+# 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 RBLNModelForImageClassificationConfig
+
+
+class RBLNViTForImageClassificationConfig(RBLNModelForImageClassificationConfig):
+    ""

optimum/rbln/transformers/models/vit/modeling_vit.py

@@ -0,0 +1,19 @@
+# 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 RBLNModelForImageClassification
+
+
+class RBLNViTForImageClassification(RBLNModelForImageClassification):
+    ""

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_wav2vec import RBLNWav2Vec2ForCTCConfig
+from .configuration_wav2vec2 import RBLNWav2Vec2ForCTCConfig
 from .modeling_wav2vec2 import RBLNWav2Vec2ForCTC

optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py

@@ -17,7 +17,7 @@ import torch
 from transformers import AutoModelForMaskedLM, Wav2Vec2ForCTC
 
 from ...modeling_generic import RBLNModelForMaskedLM
-from .configuration_wav2vec import RBLNWav2Vec2ForCTCConfig
+from .configuration_wav2vec2 import RBLNWav2Vec2ForCTCConfig
 
 
 class _Wav2Vec2(torch.nn.Module):

optimum/rbln/transformers/models/whisper/configuration_whisper.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Any, Dict
+
 import rebel
 
 from ....configuration_utils import RBLNModelConfig
@@ -29,7 +31,7 @@ class RBLNWhisperForConditionalGenerationConfig(RBLNModelConfig):
         use_attention_mask: bool = None,
         enc_max_seq_len: int = None,
         dec_max_seq_len: int = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

optimum/rbln/transformers/models/whisper/modeling_whisper.py

@@ -104,13 +104,44 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
 
 class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin):
     """
-    The Whisper Model with a language modeling head. Can be used for automatic speech recognition.
-    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    Whisper model for speech recognition and transcription optimized for RBLN NPU.
 
-    A class to convert and run pre-trained transformers based LlamaForCausalLM model on RBLN devices.
-    It implements the methods to convert a pre-trained transformers LlamaForCausalLM model into a RBLN transformer model by:
+    This model inherits from [`RBLNModel`]. It implements the methods to convert and run
+    pre-trained transformers based Whisper model on RBLN devices by:
     - transferring the checkpoint weights of the original into an optimized RBLN graph,
     - compiling the resulting graph using the RBLN compiler.
+
+    Example (Short form):
+    ```python
+    import torch
+    from transformers import AutoProcessor
+    from datasets import load_dataset
+    from optimum.rbln import RBLNWhisperForConditionalGeneration
+
+    # Load processor and dataset
+    model_id = "openai/whisper-tiny"
+    processor = AutoProcessor.from_pretrained(model_id)
+    ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+
+    # Prepare input features
+    input_features = processor(
+        ds[0]["audio"]["array"],
+        sampling_rate=ds[0]["audio"]["sampling_rate"],
+        return_tensors="pt"
+    ).input_features
+
+    # Load and compile model (or load pre-compiled model)
+    model = RBLNWhisperForConditionalGeneration.from_pretrained(
+        model_id=model_id,
+        export=True,
+        rbln_batch_size=1
+    )
+
+    # Generate transcription
+    outputs = model.generate(input_features=input_features, return_timestamps=True)
+    transcription = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+    print(f"Transcription: {transcription}")
+    ```
     """
 
     auto_model_class = AutoModelForSpeechSeq2Seq
@@ -153,11 +184,6 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
         return self.decoder
 
     def __getattr__(self, __name: str) -> Any:
-        """This is the key method to implement RBLN-Whisper.
-        Returns:
-            Any: Whisper's corresponding method
-        """
-
         def redirect(func):
             return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
 
@@ -331,12 +357,6 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
         attention_mask: Optional[torch.Tensor] = None,  # need for support transformers>=4.45.0
         **kwargs,
     ):
-        """
-        whisper don't use attention_mask,
-            attention_mask (`torch.Tensor`)`, *optional*):
-                Whisper does not support masking of the `input_features`, this argument is preserved for compatibility,
-                but it is not used. By default the silence in the input log mel spectrogram are ignored.
-        """
         return {
             "input_ids": input_ids,
             "cache_position": cache_position,

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

@@ -12,5 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_xlm_roberta import RBLNXLMRobertaModelConfig
-from .modeling_xlm_roberta import RBLNXLMRobertaModel
+from .configuration_xlm_roberta import (
+    RBLNXLMRobertaForSequenceClassificationConfig,
+    RBLNXLMRobertaModelConfig,
+)
+from .modeling_xlm_roberta import (
+    RBLNXLMRobertaForSequenceClassification,
+    RBLNXLMRobertaModel,
+)
+
+
+__all__ = [
+    "RBLNXLMRobertaModelConfig",
+    "RBLNXLMRobertaForSequenceClassificationConfig",
+    "RBLNXLMRobertaModel",
+    "RBLNXLMRobertaForSequenceClassification",
+]

optimum/rbln/transformers/models/xlm_roberta/configuration_xlm_roberta.py

@@ -12,8 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ...configuration_generic import RBLNTransformerEncoderForFeatureExtractionConfig
+from ...configuration_generic import (
+    RBLNModelForSequenceClassificationConfig,
+    RBLNTransformerEncoderForFeatureExtractionConfig,
+)
 
 
 class RBLNXLMRobertaModelConfig(RBLNTransformerEncoderForFeatureExtractionConfig):
-    pass
+    """
+    Configuration class for XLM-RoBERTa model.
+    Inherits from RBLNTransformerEncoderForFeatureExtractionConfig with no additional parameters.
+    """
+
+
+class RBLNXLMRobertaForSequenceClassificationConfig(RBLNModelForSequenceClassificationConfig):
+    """
+    Configuration class for XLM-RoBERTa sequence classification model.
+    Inherits from RBLNModelForSequenceClassificationConfig with no additional parameters.
+    """

optimum/rbln/transformers/models/xlm_roberta/modeling_xlm_roberta.py

@@ -12,9 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-from ...modeling_generic import RBLNTransformerEncoderForFeatureExtraction
+from ...modeling_generic import RBLNModelForSequenceClassification, RBLNTransformerEncoderForFeatureExtraction
 
 
 class RBLNXLMRobertaModel(RBLNTransformerEncoderForFeatureExtraction):
-    pass
+    """
+    XLM-RoBERTa base model optimized for RBLN NPU.
+    """
+
+
+class RBLNXLMRobertaForSequenceClassification(RBLNModelForSequenceClassification):
+    """
+    XLM-RoBERTa model for sequence classification tasks optimized for RBLN NPU.
+    """
+
+    rbln_model_input_names = ["input_ids", "attention_mask"]

optimum/rbln/utils/model_utils.py

@@ -12,10 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
+from typing import TYPE_CHECKING, Type
+
+
+if TYPE_CHECKING:
+    from ..modeling import RBLNModel
+
 # Prefix used for RBLN model class names
 RBLN_PREFIX = "RBLN"
 
 
+MODEL_MAPPING = {}
+
+
 def convert_hf_to_rbln_model_name(hf_model_name: str):
     """
     Convert HuggingFace model name to RBLN model name.
@@ -41,3 +51,13 @@ def convert_rbln_to_hf_model_name(rbln_model_name: str):
     """
 
     return rbln_model_name.removeprefix(RBLN_PREFIX)
+
+
+def get_rbln_model_cls(cls_name: str) -> Type["RBLNModel"]:
+    cls = getattr(importlib.import_module("optimum.rbln"), cls_name, None)
+    if cls is None:
+        if cls_name in MODEL_MAPPING:
+            cls = MODEL_MAPPING[cls_name]
+        else:
+            raise AttributeError(f"RBLNModel for {cls_name} not found.")
+    return cls

optimum/rbln/utils/submodule.py

@@ -12,19 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import importlib
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Type
 
 from transformers import PretrainedConfig
 
 from ..configuration_utils import RBLNModelConfig
+from ..utils.model_utils import get_rbln_model_cls
 
 
 if TYPE_CHECKING:
     from transformers import PreTrainedModel
 
-    from ..modeling_base import RBLNBaseModel
+    from ..modeling import RBLNModel
 
 
 class SubModulesMixin:
@@ -37,7 +37,7 @@ class SubModulesMixin:
 
     _rbln_submodules: List[Dict[str, Any]] = []
 
-    def __init__(self, *, rbln_submodules: List["RBLNBaseModel"] = [], **kwargs) -> None:
+    def __init__(self, *, rbln_submodules: List["RBLNModel"] = [], **kwargs) -> None:
         for submodule_meta, submodule in zip(self._rbln_submodules, rbln_submodules):
             setattr(self, submodule_meta["name"], submodule)
 
@@ -48,7 +48,7 @@ class SubModulesMixin:
     @classmethod
     def _export_submodules_from_model(
         cls, model: "PreTrainedModel", model_save_dir: str, rbln_config: RBLNModelConfig, **kwargs
-    ) -> List["RBLNBaseModel"]:
+    ) -> List["RBLNModel"]:
         rbln_submodules = []
         submodule_prefix = getattr(cls, "_rbln_submodule_prefix", None)
 
@@ -61,7 +61,7 @@ class SubModulesMixin:
                 torch_submodule: PreTrainedModel = getattr(model, submodule_name)
 
             cls_name = torch_submodule.__class__.__name__
-            submodule_cls: Type["RBLNBaseModel"] = getattr(importlib.import_module("optimum.rbln"), f"RBLN{cls_name}")
+            submodule_cls: Type["RBLNModel"] = get_rbln_model_cls(f"RBLN{cls_name}")
             submodule_rbln_config = getattr(rbln_config, submodule_name) or {}
 
             if isinstance(submodule_rbln_config, dict):
@@ -95,9 +95,7 @@ class SubModulesMixin:
             submodule_rbln_config = getattr(rbln_config, submodule_name)
 
             # RBLNModelConfig -> RBLNModel
-            submodule_cls: "RBLNBaseModel" = getattr(
-                importlib.import_module("optimum.rbln"), submodule_rbln_config.rbln_model_cls_name
-            )
+            submodule_cls = get_rbln_model_cls(submodule_rbln_config.rbln_model_cls_name)
 
             json_file_path = Path(model_save_dir) / submodule_name / "config.json"
             config = PretrainedConfig.from_json_file(json_file_path)

{optimum_rbln-0.8.1a0.dist-info → optimum_rbln-0.8.1a2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.1a0
+Version: 0.8.1a2
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai
@@ -28,7 +28,7 @@ Requires-Dist: packaging>=24.1
 Requires-Dist: torch==2.6.0
 Requires-Dist: torchaudio<=2.6.0
 Requires-Dist: torchvision<=0.21.0
-Requires-Dist: transformers==4.50.3
+Requires-Dist: transformers==4.51.3
 Description-Content-Type: text/markdown