optimum-rbln 0.9.3rc0__py3-none-any.whl → 0.9.4a2__py3-none-any.whl

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

@@ -13,13 +13,21 @@
 # limitations under the License.
 
 
+from typing import TYPE_CHECKING, Optional, Union
+
 import torch
-from transformers import AutoModelForMaskedLM, Wav2Vec2ForCTC
+from transformers import AutoModelForCTC, Wav2Vec2Config, Wav2Vec2ForCTC
+from transformers.modeling_outputs import CausalLMOutput
 
-from ...modeling_generic import RBLNModelForMaskedLM
+from ....configuration_utils import RBLNCompileConfig
+from ....modeling import RBLNModel
 from .configuration_wav2vec2 import RBLNWav2Vec2ForCTCConfig
 
 
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+
+
 class _Wav2Vec2(torch.nn.Module):
     def __init__(self, model: "Wav2Vec2ForCTC"):
         super().__init__()
@@ -30,13 +38,10 @@ class _Wav2Vec2(torch.nn.Module):
         return self.model.lm_head(output[0])
 
 
-class RBLNWav2Vec2ForCTC(RBLNModelForMaskedLM):
+class RBLNWav2Vec2ForCTC(RBLNModel):
     """
     Wav2Vec2 Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC).
 
-    This model inherits from [`RBLNModelForMaskedLM`]. Check the superclass documentation for the generic methods the
-    library implements for all its model.
-
     It implements the methods to convert a pre-trained Wav2Vec2 model into a RBLN Wav2Vec2 model by:
 
     - transferring the checkpoint weights of the original into an optimized RBLN graph,
@@ -44,9 +49,56 @@ class RBLNWav2Vec2ForCTC(RBLNModelForMaskedLM):
     """
 
     main_input_name = "input_values"
-    auto_model_class = AutoModelForMaskedLM
+    auto_model_class = AutoModelForCTC
     rbln_dtype = "float32"
 
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNWav2Vec2ForCTCConfig) -> torch.nn.Module:
+    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNWav2Vec2ForCTCConfig) -> torch.nn.Module:
         return _Wav2Vec2(model).eval()
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: Optional["PreTrainedModel"] = None,
+        model_config: "Wav2Vec2Config" = None,
+        rbln_config: Optional[RBLNWav2Vec2ForCTCConfig] = None,
+    ) -> RBLNWav2Vec2ForCTCConfig:
+        if rbln_config.max_seq_len is None:
+            for tokenizer in preprocessors:
+                if hasattr(tokenizer, "model_max_length"):
+                    rbln_config.max_seq_len = tokenizer.model_max_length
+                    break
+            if rbln_config.max_seq_len is None:
+                raise ValueError("`rbln_max_seq_len` should be specified!")
+
+        rbln_compile_config = RBLNCompileConfig(
+            input_info=[
+                (
+                    "input_values",
+                    [
+                        rbln_config.batch_size,
+                        rbln_config.max_seq_len,
+                    ],
+                    "float32",
+                )
+            ]
+        )
+
+        rbln_config.set_compile_cfgs([rbln_compile_config])
+        return rbln_config
+
+    def forward(
+        self, input_values: torch.Tensor, return_dict: Optional[bool] = None, **kwargs
+    ) -> Union[CausalLMOutput, tuple]:
+        """
+        Forward pass for the RBLN-optimized Wav2Vec2 model for Connectionist Temporal Classification (CTC).
+
+        Args:
+            input_values (torch.FloatTensor of shape (batch_size, sequence_length)): Float values of input raw speech waveform. Values can be obtained by loading a .flac or .wav audio file into an array of type List[float] or a numpy.ndarray, e.g. via the soundfile library (pip install soundfile). To prepare the array into input_values, the AutoProcessor should be used for padding and conversion into a tensor of type torch.FloatTensor.
+            return_dict (bool, optional): Whether or not to return a ModelOutput instead of a plain tuple.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a CausalLMOutput object.
+        """
+        return super().forward(input_values=input_values, return_dict=return_dict, **kwargs)

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

@@ -31,29 +31,63 @@ Generation utilities for Whisper.
 Modified from `transformers.models.whisper.generation_whisper.py`
 """
 
+from typing import Any, Dict, Optional, Union
+
 import torch
 import transformers
 from packaging import version
 from transformers import GenerationMixin
+from transformers.generation.configuration_utils import GenerationConfig
+from transformers.modeling_outputs import ModelOutput
 from transformers.models.whisper.generation_whisper import WhisperGenerationMixin
 
 
 class RBLNWhisperGenerationMixin(WhisperGenerationMixin, GenerationMixin):
-    def generate(self, *args, generation_config=None, **kwargs):
-        num_beams = kwargs.get(
-            "num_beams",
-            generation_config.num_beams
-            if hasattr(generation_config, "num_beams") and generation_config.num_beams is not None
-            else 1,
-        )
-        if num_beams > 1:
-            raise ValueError(
-                f"Beam search is not supported in RBLNWhisperGenerationMixin. "
-                f"Received num_beams={num_beams}, but only num_beams=1 is allowed. "
-                f"Please set num_beams=1 for greedy search or adjust your configuration."
-            )
+    def generate(
+        self,
+        input_features: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
+        return_segments: Optional[bool] = None,
+        return_timestamps: Optional[bool] = None,
+        return_token_timestamps: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[ModelOutput, Dict[str, Any], torch.LongTensor]:
+        """
+        The generate function is utilized in its standard form as in the HuggingFace transformers library. User can use this function to generate text from the model.
+        Check the [HuggingFace transformers documentation](https://huggingface.co/docs/transformers/v4.57.1/en/model_doc/whisper#transformers.WhisperForConditionalGeneration.generate) for more details.
+
+        Args:
+            input_features(torch.Tensor, optional): The input features to the model.
+            attention_mask(torch.Tensor, optional): Attention mask needs to be passed when doing long-form transcription using a batch size > 1.
+            generation_config(GenerationConfig, optional): The generation configuration to be used as base parametrization for the generation call. **kwargs passed to generate matching the attributes of generation_config will override them.
+                If generation_config is not provided, the default will be used, which had the following loading priority: 1) from the generation_config.json model file, if it exists; 2) from the model configuration.
+                Please note that unspecified parameters will inherit [GenerationConfig](https://huggingface.co/docs/transformers/v4.57.1/en/main_classes/text_generation#transformers.GenerationConfig)’s default values.
+            return_segments(bool, optional): Whether to return segments.
+            return_timestamps(bool, optional): Whether to return the timestamps with the text. For audios longer than 30 seconds, it is necessary to set return_timestamps=True.
+            return_token_timestamps(bool, optional): Whether to return token timestamps.
+            kwargs(dict[str, Any], optional): Additional arguments passed to the generate function.
+
+        Returns:
+            Transcribes or translates log-mel input features to a sequence of auto-regressively generated token ids.
+        """
+        if kwargs.get("num_beams", None) is not None:
+            if kwargs.get("num_beams") != 1:
+                raise ValueError(
+                    "Beam search is not supported in RBLNWhisperGenerationMixin. "
+                    "Received num_beams={num_beams}, but only num_beams=1 is allowed. "
+                    "Please set num_beams=1 for greedy search or adjust your configuration."
+                )
 
-        return super().generate(*args, **kwargs)
+        return super().generate(
+            input_features,
+            attention_mask=attention_mask,
+            generation_config=generation_config,
+            return_segments=return_segments,
+            return_timestamps=return_timestamps,
+            return_token_timestamps=return_token_timestamps,
+            **kwargs,
+        )
 
     def _postprocess_outputs(
         self,

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

@@ -203,7 +203,7 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
         raise NotImplementedError
 
     @classmethod
-    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNWhisperForConditionalGenerationConfig):
+    def _wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNWhisperForConditionalGenerationConfig):
         return WhisperWrapper(
             model,
             use_attention_mask=rbln_config.use_attention_mask,
@@ -213,7 +213,7 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
     @classmethod
     @torch.inference_mode()
     def get_compiled_model(cls, model, rbln_config: RBLNWhisperForConditionalGenerationConfig):
-        wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
+        wrapped_model = cls._wrap_model_if_needed(model, rbln_config)
 
         enc_compile_config = rbln_config.compile_cfgs[0]
         dec_compile_config = rbln_config.compile_cfgs[1]

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

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional, Union
+
+import torch
+from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions, SequenceClassifierOutput
+
 from ...modeling_generic import RBLNModelForSequenceClassification, RBLNTransformerEncoderForFeatureExtraction
 
 
@@ -20,6 +25,30 @@ class RBLNXLMRobertaModel(RBLNTransformerEncoderForFeatureExtraction):
     XLM-RoBERTa base model optimized for RBLN NPU.
     """
 
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[BaseModelOutputWithPoolingAndCrossAttentions, tuple]:
+        """
+        Forward pass for the RBLN-optimized XLM-RoBERTa base model.
+
+        Args:
+            input_ids (torch.Tensor of shape (batch_size, sequence_length), optional): Indices of input sequence tokens in the vocabulary.
+            attention_mask (torch.Tensor of shape (batch_size, sequence_length), optional): Mask to avoid performing attention on padding token indices.
+            token_type_ids (torch.Tensor of shape (batch_size, sequence_length), optional): Segment token indices to indicate different portions of the inputs.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a BaseModelOutputWithPoolingAndCrossAttentions object.
+        """
+
+        if token_type_ids is not None:
+            kwargs.setdefault("token_type_ids", token_type_ids)
+
+        return super().forward(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
+
 
 class RBLNXLMRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     """
@@ -27,3 +56,27 @@ class RBLNXLMRobertaForSequenceClassification(RBLNModelForSequenceClassification
     """
 
     rbln_model_input_names = ["input_ids", "attention_mask"]
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[SequenceClassifierOutput, tuple]:
+        """
+        Forward pass for the RBLN-optimized XLM-RoBERTa model for sequence classification.
+
+        Args:
+            input_ids (torch.LongTensor of shape (batch_size, sequence_length), optional): Indices of input sequence tokens in the vocabulary.
+            attention_mask (torch.FloatTensor of shape (batch_size, sequence_length), optional): Mask to avoid performing attention on padding token indices.
+            token_type_ids (torch.LongTensor of shape (batch_size, sequence_length), optional): Segment token indices to indicate first and second portions of the inputs.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a SequenceClassifierOutput object.
+        """
+
+        if token_type_ids is not None:
+            kwargs.setdefault("token_type_ids", token_type_ids)
+
+        return super().forward(input_ids=input_ids, attention_mask=attention_mask, **kwargs)

optimum/rbln/transformers/utils/rbln_quantization.py

@@ -123,6 +123,15 @@ class RBLNQuantizationConfig(RBLNSerializableConfigProtocol):
         if self.RBLN_QUANT_BITS_ENV in os.environ:
             os.environ.pop(self.RBLN_QUANT_BITS_ENV)
 
+    @property
+    def nbits_per_param(self) -> int:
+        if self.weights in ["int4", "fp4"]:
+            return 4
+        elif self.weights in ["int8", "fp8"]:
+            return 8
+        else:
+            raise ValueError(f"Invalid weights: {self.weights}")
+
 
 class QuantizedLayerFactory:
     def __init__(self, quantization_config: RBLNQuantizationConfig):

optimum/rbln/utils/deprecation.py

@@ -0,0 +1,213 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+# 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.
+
+# **********************************************************************************
+# * NOTE: This file has been modified from its original version in              *
+# * the Hugging Face transformers library.                                      *
+# * Original source:                                                            *
+# * https://github.com/huggingface/transformers/blob/v4.57.1/src/transformers/utils/deprecation.py
+# **********************************************************************************
+
+import inspect
+from enum import Enum
+from functools import wraps
+from typing import Callable, Optional
+
+import packaging.version
+
+from ..__version__ import __version__
+from .logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+def warn_deprecated_npu(npu: Optional[str] = None):
+    import rebel
+
+    npu = npu or rebel.get_npu_name()
+    if npu == "RBLN-CA02":
+        logger.warning_once(
+            "Support for the RBLN-CA02 device is provided only up to optimum-rbln v0.8.0 and has reached end of life.",
+        )
+
+
+class Action(Enum):
+    NONE = "none"
+    NOTIFY = "notify"
+    RAISE = "raise"
+
+
+# Scenario Table for Deprecation Strategy Example
+# Assume that current version is v0.9.6 and the deprecated version is v0.10.0
+# |--------------------|----------------|----------------|---------------------------------------------|--------------------------------------------------------------------------------------|----------------------------------------------------------------------|
+# | Type               | v0.9.6 (as_is) | v0.9.6 (to_be) | v0.9.6 Patch                                | v0.9.7 Action                                                                        | v0.10.0+ Action                                                      |
+# |--------------------|----------------|----------------|---------------------------------------------|--------------------------------------------------------------------------------------|----------------------------------------------------------------------|
+# | Modify (Key Name)  | a: bool        | a': bool       | Add a', Keep a                              | 1. Only 'a' provided: replace a -> a' & future warning                               | In v0.10.0, raise error once, then remove decorator.                 |
+# |                    |                |                |                                             | 2. Both 'a' & 'a'' provided: ignore 'a' value & future warning                       |                                                                      |
+# |--------------------|----------------|----------------|---------------------------------------------|--------------------------------------------------------------------------------------|----------------------------------------------------------------------|
+# | Modify (Value Type)| b: bool        | b: str         | b: Union[str, bool]                         | 'bool' value provided for 'b': replace with corresponding 'str' & future warning     | In v0.10.0, raise error once, then remove decorator.                 |
+# |                    |                |                |                                             |                                                                                      |                                                                      |
+# |--------------------|----------------|----------------|---------------------------------------------|--------------------------------------------------------------------------------------|----------------------------------------------------------------------|
+# | Deletion           | c              | -              | Delete c or Keep c (flexible)               | ignore c & future warning                                                            | In v0.10.0, raise error once, then remove decorator.                 |
+# |                    |                |                |                                             |                                                                                      |                                                                      |
+# |--------------------|----------------|----------------|---------------------------------------------|--------------------------------------------------------------------------------------|----------------------------------------------------------------------|
+# | Addition           | -              | d              | Add d, set default_value for d              | No action needed as default value is set                                             | Keep default value                                                   |
+# |--------------------|----------------|----------------|---------------------------------------------|--------------------------------------------------------------------------------------|----------------------------------------------------------------------|
+
+
+def deprecate_kwarg(
+    old_name: str,
+    version: str,
+    new_name: Optional[str] = None,
+    deprecated_type: Optional[type] = None,
+    value_replacer: Optional[Callable] = None,
+    raise_if_greater_or_equal_version: bool = True,
+    raise_if_both_names: bool = False,
+    additional_message: Optional[str] = None,
+):
+    """
+    Function or method decorator to notify users about deprecated keyword arguments, replacing them with a new name if specified,
+    or handling deprecated value types.
+
+    This decorator allows you to:
+    - Notify users when a keyword argument name is deprecated (Scenario 'a', 'c').
+    - Notify users when a specific value type for an argument is deprecated (Scenario 'b').
+    - Automatically replace deprecated keyword arguments with new ones.
+    - Automatically replace deprecated values with new ones using a replacer function.
+    - Raise an error if deprecated arguments are used, depending on the specified conditions.
+
+    By default, the decorator notifies the user about the deprecated argument while the `optimum.rbln.__version__` < specified `version`
+    in the decorator. To keep notifications with any version `warn_if_greater_or_equal_version=True` can be set.
+
+    Parameters:
+        old_name (`str`):
+            Name of the deprecated keyword argument, or the argument with a deprecated value type.
+        version (`str`):
+            The version in which the keyword argument or value type was (or will be) deprecated.
+        new_name (`Optional[str]`, *optional*):
+            The new name for the deprecated keyword argument. If specified, the deprecated keyword argument will be replaced with this new name (Scenario 'a').
+        deprecated_type (`type`, *optional*):
+            The deprecated type for the keyword argument specified by `old_name` (Scenario 'b').
+            If this is set, `new_name` should typically be `None`.
+        value_replacer (`Callable`, *optional*):
+            A function that takes the old (deprecated type) value and returns a new value (Scenario 'b').
+            Used in conjunction with `deprecated_type`. If provided, the value will be automatically converted.
+        raise_if_greater_or_equal_version (`bool`, *optional*, defaults to `False`):
+            Whether to raise `ValueError` if current `optimum.rbln.` version is greater or equal to the deprecated version.
+        raise_if_both_names (`bool`, *optional*, defaults to `False`):
+            Whether to raise `ValueError` if both deprecated and new keyword arguments are set (only for Scenario 'a').
+        additional_message (`Optional[str]`, *optional*):
+            An additional message to append to the default deprecation message.
+
+    Raises:
+        ValueError:
+            If raise_if_greater_or_equal_version is True and the current version is greater than or equal to the deprecated version, or if raise_if_both_names is True and both old and new keyword arguments are provided.
+
+    Returns:
+        Callable:
+            A wrapped function that handles the deprecated keyword arguments according to the specified parameters.
+    """
+
+    deprecated_version = packaging.version.parse(version)
+    current_version = packaging.version.parse(__version__)
+    is_greater_or_equal_version = current_version >= deprecated_version
+
+    if is_greater_or_equal_version:
+        version_message = f"and removed starting from version {version}"
+    else:
+        version_message = f"and will be removed in version {version}"
+
+    def wrapper(func):
+        # Required for better warning message
+        sig = inspect.signature(func)
+        function_named_args = set(sig.parameters.keys())
+        is_instance_method = "self" in function_named_args
+        is_class_method = "cls" in function_named_args
+
+        @wraps(func)
+        def wrapped_func(*args, **kwargs):
+            # Get class + function name (just for better warning message)
+            func_name = func.__name__
+            if is_instance_method:
+                func_name = f"{args[0].__class__.__name__}.{func_name}"
+            elif is_class_method:
+                func_name = f"{args[0].__name__}.{func_name}"
+
+            minimum_action = Action.NONE
+            message = None
+
+            # Scenario A: Rename (e.g., a -> a')
+            if new_name is not None:
+                if old_name in kwargs and new_name in kwargs:
+                    minimum_action = Action.RAISE if raise_if_both_names else Action.NOTIFY
+                    message = f"Both `{old_name}` and `{new_name}` are set for `{func_name}`. Using `{new_name}={kwargs[new_name]}` and ignoring deprecated `{old_name}={kwargs[old_name]}`."
+                    kwargs.pop(old_name)
+
+                elif old_name in kwargs and new_name not in kwargs:
+                    minimum_action = Action.NOTIFY
+                    message = (
+                        f"`{old_name}` is deprecated {version_message} for `{func_name}`. Use `{new_name}` instead."
+                    )
+                    kwargs[new_name] = kwargs.pop(old_name)
+
+            # Scenario B: Value Type Change (e.g., b: bool -> str)
+            if deprecated_type is not None:
+                key_to_check = old_name if new_name is None else new_name  # For Senario A + B Mixed
+                if key_to_check in kwargs and isinstance(kwargs[key_to_check], deprecated_type):
+                    minimum_action = Action.NOTIFY
+                    old_value = kwargs[key_to_check]
+                    message = f"Using type `{deprecated_type.__name__}` for argument `{key_to_check}` in `{func_name}` is deprecated {version_message}."
+
+                    if value_replacer:
+                        try:
+                            new_value = value_replacer(old_value)
+                            kwargs[key_to_check] = new_value
+                            message += f" Value `{old_value}` has been automatically replaced with `{new_value}`."
+                        except Exception as e:
+                            logger.error(f"Error during deprecated value replacement for {key_to_check}: {e}")
+                            message += f" Automatic replacement failed: {e}. Passing original value."
+                    else:
+                        raise ValueError(
+                            f"value_replacer should be provided when deprecated_type is set for {key_to_check} in {func_name}"
+                        )
+
+            # Scenario C: Deletion (e.g., c)
+            if old_name in kwargs and new_name is None and deprecated_type is None:
+                minimum_action = Action.NOTIFY
+                message = f"`{old_name}` is deprecated {version_message} for `{func_name}`."
+                kwargs.pop(old_name)
+
+            if message is not None and additional_message is not None:
+                message = f"{message} {additional_message}"
+
+            # update minimum_action if argument is ALREADY deprecated (current version >= deprecated version)
+            if is_greater_or_equal_version:
+                # change to NOTIFY -> RAISE  in case we want to raise error for already deprecated arguments
+                if raise_if_greater_or_equal_version:
+                    minimum_action = Action.RAISE
+
+            # raise error or notify user
+            if minimum_action == Action.RAISE:
+                raise ValueError(message)
+            elif minimum_action == Action.NOTIFY:
+                # DeprecationWarning is ignored by default, so we use FutureWarning instead
+                logger.warning(message, stacklevel=2)
+
+            return func(*args, **kwargs)
+
+        return wrapped_func
+
+    return wrapper

optimum/rbln/utils/hub.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 from pathlib import Path
 from typing import List, Optional, Union
 
@@ -67,15 +68,25 @@ def validate_files(
     location: str,
 ):
     """Validate the presence and count of required files."""
-    if len(files) == 0:
-        raise FileNotFoundError(f"Could not find any rbln model file in {location}")
-
     if len(config_files) == 0:
         raise FileNotFoundError(f"Could not find `rbln_config.json` file in {location}")
 
     if len(config_files) > 1:
         raise FileExistsError(f"Multiple rbln_config.json files found in {location}. This is not expected.")
 
+    try:
+        with open(config_files[0], "r") as f:
+            config_data = json.load(f)
+        compile_cfgs = config_data.get("_compile_cfgs", [])
+        if len(compile_cfgs) == 0:
+            # If compile_cfgs is empty, we don't need .rbln files
+            return
+    except (json.JSONDecodeError, KeyError, OSError):
+        pass
+
+    if len(files) == 0:
+        raise FileNotFoundError(f"Could not find any rbln model file in {location}")
+
 
 def _get_huggingface_token(token: Union[bool, str]) -> str:
     if isinstance(token, str):

optimum/rbln/utils/import_utils.py

@@ -142,7 +142,11 @@ def check_version_compats() -> None:
         try:
             dep_version = importlib.metadata.version(compat.package_name)
         except importlib.metadata.PackageNotFoundError:
-            warnings.warn(f"optimum-rbln requires {compat.package_name} to be installed.", ImportWarning)
+            warnings.warn(
+                f"optimum-rbln requires {compat.package_name} to be installed.",
+                ImportWarning,
+                stacklevel=2,
+            )
             continue
         # For versions 0.7.2 and above, don't show warning for rebel-compiler if base versions match
 
@@ -160,6 +164,7 @@ def check_version_compats() -> None:
                     f"For optimal performance and compatibility, please ensure both packages share the same major and minor version numbers. "
                     "Please refer to our SDK release notes at https://docs.rbln.ai/about_atom/release_note.html",
                     ImportWarning,
+                    stacklevel=2,
                 )
         else:
             if not Version(compat.min_version) <= Version(dep_version) < Version(compat.max_version):
@@ -167,4 +172,5 @@ def check_version_compats() -> None:
                     f"optimum-rbln v{my_version} is compatible to {compat.package_name} v{compat.min_version} to v{compat.max_version}. (you are currently using v{dep_version})\n"
                     "Please refer to our SDK release notes at https://docs.rbln.ai/about_atom/release_note.html",
                     ImportWarning,
+                    stacklevel=2,
                 )

optimum/rbln/utils/runtime_utils.py

@@ -20,6 +20,38 @@ import rebel
 import torch
 
 
+def get_available_dram(npu: Optional[str] = None) -> int:
+    """
+    Get the available DRAM size of the specified NPU.
+
+    Args:
+        npu : Optional[str], default=None
+            The NPU to get the available DRAM size.
+            If None, the function will attempt to retrieve through `ensure_valid_npu()`
+
+    Returns:
+        int
+            The available DRAM size in bytes.
+    """
+    if npu is None:
+        if not rebel.npu_is_available(0):
+            raise RuntimeError("No NPU is available to get available DRAM size.")
+
+        npu = rebel.get_npu_name(0)
+
+    if npu.startswith("RBLN-CR"):
+        # TODO(jongho): Assuming 4 chiplets.
+        DRAM_NBYTES = 144 * 2**30
+        SYS_DRAM_NBYTES = 4 * 2**30
+    elif npu.startswith("RBLN-CA"):
+        DRAM_NBYTES = 16 * 2**30
+        SYS_DRAM_NBYTES = 288 * 2**20
+    else:
+        raise ValueError(f"Unknown npu name: {npu}")
+
+    return DRAM_NBYTES - SYS_DRAM_NBYTES
+
+
 def normalize_npu(npu: str) -> str:
     """Normalize the NPU string by removing the form factor."""
     match = re.match(r"(RBLN-CA|RBLN-CR)(\d+)", npu)

optimum/rbln/utils/submodule.py

@@ -37,7 +37,9 @@ class SubModulesMixin:
 
     _rbln_submodules: List[Dict[str, Any]] = []
 
-    def __init__(self, *, rbln_submodules: List["RBLNModel"] = [], **kwargs) -> None:
+    def __init__(self, *, rbln_submodules: Optional[List["RBLNModel"]] = None, **kwargs) -> None:
+        if rbln_submodules is None:
+            rbln_submodules = []
         for submodule_meta, submodule in zip(self._rbln_submodules, rbln_submodules):
             setattr(self, submodule_meta["name"], submodule)
 

{optimum_rbln-0.9.3rc0.dist-info → optimum_rbln-0.9.4a2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.9.3rc0
+Version: 0.9.4a2
 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
@@ -24,7 +24,7 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: <3.14,>=3.9
 Requires-Dist: accelerate>=1.0.1
-Requires-Dist: diffusers==0.35.1
+Requires-Dist: diffusers==0.35.2
 Requires-Dist: packaging>=24.1
 Requires-Dist: torch==2.8.0
 Requires-Dist: torchaudio<=2.8.0