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

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.0.post2.dist-info → optimum_rbln-0.8.1a1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.0.post2
+Version: 0.8.1a1
 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