optimum-rbln 0.9.3__py3-none-any.whl → 0.9.3rc0__py3-none-any.whl

optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py

@@ -20,9 +20,8 @@ import rebel
 import torch
 from rebel.compile_context import CompileContext
 from transformers import AutoModelForSeq2SeqLM, PretrainedConfig, PreTrainedModel
-from transformers.generation.configuration_utils import GenerationConfig
 from transformers.generation.utils import GenerationMixin
-from transformers.modeling_outputs import BaseModelOutput, ModelOutput, Seq2SeqLMOutput
+from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
 
 from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
@@ -34,7 +33,7 @@ from .configuration_seq2seq import RBLNModelForSeq2SeqLMConfig
 logger = get_logger(__name__)
 
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, GenerationConfig, PretrainedConfig
 
 
 class RBLNRuntimeEncoder(RBLNPytorchRuntime):
@@ -141,7 +140,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, GenerationMixin, ABC):
     @classmethod
     @torch.inference_mode()
     def get_compiled_model(cls, model: PreTrainedModel, rbln_config: RBLNModelForSeq2SeqLMConfig):
-        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]
@@ -222,6 +221,12 @@ class RBLNModelForSeq2SeqLM(RBLNModel, GenerationMixin, ABC):
             model_config, "max_position_embeddings", None
         )
 
+        pad_token_id = getattr(model_config, "pad_token_id", None)
+        pad_token_id = pad_token_id or getattr(model_config, "bos_token_id", None)
+        pad_token_id = pad_token_id or getattr(model_config, "eos_token_id", None)
+        pad_token_id = pad_token_id or -1
+        rbln_config.pad_token_id = pad_token_id
+
         if rbln_config.enc_max_seq_len is None:
             enc_max_seq_len = max_position_embeddings
             for tokenizer in preprocessors:
@@ -427,7 +432,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, GenerationMixin, ABC):
         inputs_tensor = torch.nn.functional.pad(
             inputs_tensor,
             (0, self.rbln_config.enc_max_seq_len - input_len),
-            value=self.config.pad_token_id,
+            value=self.rbln_config.pad_token_id,
         )
         model_kwargs["attention_mask"] = torch.nn.functional.pad(
             model_kwargs["attention_mask"], (0, self.rbln_config.enc_max_seq_len - input_len)
@@ -446,32 +451,3 @@ class RBLNModelForSeq2SeqLM(RBLNModel, GenerationMixin, ABC):
             model_kwargs["encoder_outputs"] = encoder(**encoder_kwargs, block_tables=block_tables)
 
         return model_kwargs
-
-    def generate(
-        self,
-        input_ids: torch.LongTensor,
-        attention_mask: Optional[torch.LongTensor] = None,
-        generation_config: Optional[GenerationConfig] = None,
-        **kwargs,
-    ) -> Union[ModelOutput, 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/main_classes/text_generation#transformers.GenerationMixin.generate) for more details.
-
-        Args:
-            input_ids (torch.LongTensor): The input ids to the model.
-            attention_mask (torch.LongTensor, optional): The attention mask to the model.
-            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.
-            kwargs (dict[str, Any], optional): Additional arguments passed to the generate function. See the HuggingFace transformers documentation for more details.
-
-        Returns:
-            Generates sequences of token ids for models with a language modeling head.
-        """
-        if generation_config is not None:
-            kwargs["generation_config"] = generation_config
-        if attention_mask is not None:
-            kwargs["attention_mask"] = attention_mask
-
-        return super().generate(input_ids, **kwargs)

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

@@ -66,9 +66,7 @@ class RBLNSiglipVisionModel(RBLNModel):
     _tp_support = False
 
     @classmethod
-    def _wrap_model_if_needed(
-        cls, model: torch.nn.Module, rbln_config: RBLNSiglipVisionModelConfig
-    ) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNSiglipVisionModelConfig) -> torch.nn.Module:
         wrapper_cfg = {
             "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
             "output_hidden_states": rbln_config.output_hidden_states,
@@ -124,20 +122,6 @@ class RBLNSiglipVisionModel(RBLNModel):
         interpolate_pos_encoding: bool = False,
         **kwargs: Any,
     ) -> Union[Tuple, BaseModelOutputWithPooling]:
-        """
-        Forward pass for the RBLN-optimized SigLIP vision model.
-
-        Args:
-            pixel_values (torch.FloatTensor of shape (batch_size, num_channels, image_size, image_size), optional): The tensors corresponding to the input images. Pixel values can be obtained using ViTImageProcessor. See ViTImageProcessor.call() for details (processor_class uses ViTImageProcessor for processing images).
-            return_dict (bool, optional): Whether or not to return a ModelOutput instead of a plain tuple.
-            output_attentions (bool, optional): Whether or not to return the attentions tensors of all attention layers. See attentions under returned tensors for more detail.
-            output_hidden_states (bool, optional): Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for more detail.
-            interpolate_pos_encoding (bool, defaults to False): Whether to interpolate the pre-trained position encodings.
-
-        Returns:
-            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a BaseModelOutputWithPooling object.
-        """
-
         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

optimum/rbln/transformers/models/swin/modeling_swin.py

@@ -203,7 +203,7 @@ class _SwinBackbone(torch.nn.Module):
 
 class RBLNSwinBackbone(RBLNModel):
     @classmethod
-    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNSwinBackboneConfig) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNSwinBackboneConfig) -> torch.nn.Module:
         for layer in model.encoder.layers:
             for block in layer.blocks:
                 block.get_attn_mask = types.MethodType(get_attn_mask, block)
@@ -278,19 +278,6 @@ class RBLNSwinBackbone(RBLNModel):
         output_hidden_states: bool = None,
         **kwargs,
     ) -> Union[Tuple, BackboneOutput]:
-        """
-        Forward pass for the RBLN-optimized Swin backbone model.
-
-        Args:
-            pixel_values (torch.FloatTensor of shape (batch_size, num_channels, image_size, image_size), optional): The tensors corresponding to the input images. Pixel values can be obtained using ViTImageProcessor. See ViTImageProcessor.call() for details (processor_class uses ViTImageProcessor for processing images).
-            return_dict (bool, optional): Whether or not to return a ModelOutput instead of a plain tuple.
-            output_attentions (bool, optional): Whether or not to return the attentions tensors of all attention layers. See attentions under returned tensors for more detail.
-            output_hidden_states (bool, optional): Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for more detail.
-
-        Returns:
-            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a BackboneOutput object.
-        """
-
         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__}."

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

@@ -68,7 +68,7 @@ class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
     output_class = BaseModelOutputWithPastAndCrossAttentions
 
     @classmethod
-    def _wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNT5EncoderModelConfig):
+    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNT5EncoderModelConfig):
         return T5EncoderWrapper(model)
 
     @classmethod
@@ -113,7 +113,7 @@ class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
     support_causal_attn = False
 
     @classmethod
-    def _wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNT5ForConditionalGenerationConfig):
+    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNT5ForConditionalGenerationConfig):
         return T5Wrapper(
             model, enc_max_seq_len=rbln_config.enc_max_seq_len, dec_max_seq_len=rbln_config.dec_max_seq_len
         )

optimum/rbln/transformers/models/time_series_transformer/modeling_time_series_transformer.py

@@ -153,7 +153,7 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
             return redirect(val)
 
     @classmethod
-    def _wrap_model_if_needed(
+    def wrap_model_if_needed(
         self, model: "PreTrainedModel", rbln_config: RBLNTimeSeriesTransformerForPredictionConfig
     ):
         return TimeSeriesTransformersWrapper(model, rbln_config.num_parallel_samples)
@@ -161,7 +161,7 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
     @classmethod
     @torch.inference_mode()
     def get_compiled_model(cls, model, rbln_config: RBLNTimeSeriesTransformerForPredictionConfig):
-        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]
@@ -353,20 +353,6 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
         static_real_features: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> SampleTSPredictionOutput:
-        """
-        Generate pass for the RBLN-optimized Time Series Transformer model for time series forecasting.
-
-        Args:
-            past_values (torch.FloatTensor of shape (batch_size, sequence_length) or (batch_size, sequence_length, input_size)): Past values of the time series, that serve as context in order to predict the future.
-            past_time_features (torch.FloatTensor of shape (batch_size, sequence_length, num_features)): Required time features, which the model internally will add to past_values.
-            future_time_features (torch.FloatTensor of shape (batch_size, prediction_length, num_features)): Required time features for the prediction window, which the model internally will add to future_values.
-            past_observed_mask (torch.BoolTensor of shape (batch_size, sequence_length) or (batch_size, sequence_length, input_size), optional): Boolean mask to indicate which past_values were observed and which were missing.
-            static_categorical_features (torch.LongTensor of shape (batch_size, number of static categorical features), optional): Optional static categorical features for which the model will learn an embedding, which it will add to the values of the time series.
-            static_real_features (torch.FloatTensor of shape (batch_size, number of static real features), optional): Optional static real features which the model will add to the values of the time series.
-
-        Returns:
-            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a SampleTSPredictionOutput object.
-        """
         self.validate_batch_size(**{k: v for k, v in locals().items() if isinstance(v, torch.Tensor)})
 
         outputs = self.encoder(

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

@@ -12,11 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Tuple, Union
-
-import torch
-from transformers.modeling_outputs import ImageClassifierOutput
-
 from ...modeling_generic import RBLNModelForImageClassification
 
 
@@ -28,17 +23,3 @@ class RBLNViTForImageClassification(RBLNModelForImageClassification):
     on RBLN devices, supporting image classification with transformer-based architectures
     that process images as sequences of patches.
     """
-
-    def forward(self, pixel_values: torch.Tensor, **kwargs) -> Union[ImageClassifierOutput, Tuple]:
-        """
-        Forward pass for the RBLN-optimized Vision Transformer model for image classification.
-
-        Args:
-            pixel_values (torch.FloatTensor of shape (batch_size, channels, height, width)):
-                The tensors corresponding to the input images.
-
-        Returns:
-            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns an ImageClassifierOutput object.
-
-        """
-        return super().forward(pixel_values, **kwargs)

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

@@ -12,12 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Optional
+from ...configuration_generic import RBLNModelForMaskedLMConfig
 
-from ....configuration_utils import RBLNModelConfig
 
-
-class RBLNWav2Vec2ForCTCConfig(RBLNModelConfig):
+class RBLNWav2Vec2ForCTCConfig(RBLNModelForMaskedLMConfig):
     """
     Configuration class for RBLNWav2Vec2ForCTC.
 
@@ -25,14 +23,4 @@ class RBLNWav2Vec2ForCTCConfig(RBLNModelConfig):
     RBLN-optimized Wav2Vec2 models for Connectionist Temporal Classification (CTC) tasks.
     """
 
-    def __init__(
-        self,
-        max_seq_len: Optional[int] = None,
-        batch_size: Optional[int] = None,
-        **kwargs: Any,
-    ):
-        super().__init__(**kwargs)
-        self.max_seq_len = max_seq_len
-        self.batch_size = batch_size or 1
-        if not isinstance(self.batch_size, int) or self.batch_size < 0:
-            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+    rbln_model_input_names = ["input_values"]

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

@@ -13,21 +13,13 @@
 # limitations under the License.
 
 
-from typing import TYPE_CHECKING, Optional, Union
-
 import torch
-from transformers import AutoModelForCTC, Wav2Vec2Config, Wav2Vec2ForCTC
-from transformers.modeling_outputs import CausalLMOutput
+from transformers import AutoModelForMaskedLM, Wav2Vec2ForCTC
 
-from ....configuration_utils import RBLNCompileConfig
-from ....modeling import RBLNModel
+from ...modeling_generic import RBLNModelForMaskedLM
 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__()
@@ -38,10 +30,13 @@ class _Wav2Vec2(torch.nn.Module):
         return self.model.lm_head(output[0])
 
 
-class RBLNWav2Vec2ForCTC(RBLNModel):
+class RBLNWav2Vec2ForCTC(RBLNModelForMaskedLM):
     """
     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,
@@ -49,56 +44,9 @@ class RBLNWav2Vec2ForCTC(RBLNModel):
     """
 
     main_input_name = "input_values"
-    auto_model_class = AutoModelForCTC
+    auto_model_class = AutoModelForMaskedLM
     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,63 +31,29 @@ 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,
-        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(
-            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 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."
+            )
+
+        return super().generate(*args, **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,11 +12,6 @@
 # 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
 
 
@@ -25,25 +20,6 @@ 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,
-        **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.
-
-        Returns:
-            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a BaseModelOutputWithPoolingAndCrossAttentions object.
-        """
-
-        return super().forward(input_ids, attention_mask, **kwargs)
-
 
 class RBLNXLMRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     """
@@ -51,22 +27,3 @@ 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,
-        **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.
-
-        Returns:
-            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a SequenceClassifierOutput object.
-        """
-
-        return super().forward(input_ids, attention_mask, **kwargs)

optimum/rbln/transformers/utils/rbln_quantization.py

@@ -123,15 +123,6 @@ 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/depreacate_utils.py

@@ -0,0 +1,16 @@
+from typing import Optional
+
+import rebel
+
+from .logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+def warn_deprecated_npu(npu: Optional[str] = None):
+    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.",
+        )

optimum/rbln/utils/hub.py

@@ -12,7 +12,6 @@
 # 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
 
@@ -68,25 +67,15 @@ 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/runtime_utils.py

@@ -20,38 +20,6 @@ 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-0.9.3.dist-info → optimum_rbln-0.9.3rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.9.3
+Version: 0.9.3rc0
 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.2
+Requires-Dist: diffusers==0.35.1
 Requires-Dist: packaging>=24.1
 Requires-Dist: torch==2.8.0
 Requires-Dist: torchaudio<=2.8.0