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

optimum/rbln/transformers/models/distilbert/modeling_distilbert.py

@@ -12,11 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple, Union
-
-import torch
-from transformers.modeling_outputs import QuestionAnsweringModelOutput
-
 from ...modeling_generic import RBLNModelForQuestionAnswering
 
 
@@ -30,22 +25,3 @@ class RBLNDistilBertForQuestionAnswering(RBLNModelForQuestionAnswering):
     """
 
     rbln_model_input_names = ["input_ids", "attention_mask"]
-
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        **kwargs,
-    ) -> Union[Tuple, QuestionAnsweringModelOutput]:
-        """
-        Forward pass for the RBLN-optimized DistilBERT model for question answering tasks.
-
-        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 QuestionAnsweringModelOutput object.
-        """
-
-        return super().forward(input_ids, attention_mask, **kwargs)

optimum/rbln/transformers/models/dpt/modeling_dpt.py

@@ -13,11 +13,6 @@
 # limitations under the License.
 
 
-from typing import Tuple, Union
-
-import torch
-from transformers.modeling_outputs import DepthEstimatorOutput
-
 from ...modeling_generic import RBLNModelForDepthEstimation
 
 
@@ -28,15 +23,3 @@ class RBLNDPTForDepthEstimation(RBLNModelForDepthEstimation):
     This class provides hardware-accelerated inference for DPT (Dense Prediction Transformer)
     models on RBLN devices, supporting monocular depth estimation from single images.
     """
-
-    def forward(self, pixel_values: torch.Tensor, **kwargs) -> Union[Tuple, DepthEstimatorOutput]:
-        """
-        Forward pass for the RBLN-optimized DPT model.
-
-        Args:
-            pixel_values (torch.FloatTensor of shape (batch_size, num_channels, image_size, image_size)): The tensors corresponding to the input images.
-
-        Returns:
-            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a DepthEstimatorOutput object.
-        """
-        return super().forward(pixel_values, **kwargs)

optimum/rbln/transformers/models/gemma3/modeling_gemma3.py

@@ -99,7 +99,9 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMix
         return True
 
     @classmethod
-    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
+    def get_pytorch_model(cls, *args, **kwargs):
+        model = super().get_pytorch_model(*args, **kwargs)
+
         with no_init_weights():
             model_cls_name = model.model.language_model.__class__.__name__
             causal_model_cls_name = model_cls_name.replace("TextModel", "ForCausalLM")
@@ -133,7 +135,7 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMix
         return self.language_model.get_input_embeddings()
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
+    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         return model.multi_modal_projector
 
     @classmethod
@@ -480,7 +482,7 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     @classmethod
     @torch.inference_mode()
     def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNGemma3ForCausalLMConfig):
-        wrapped_model = cls._wrap_model_if_needed(model, rbln_config)
+        wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
 
         rbln_compile_configs = rbln_config.compile_cfgs
         prefill_compile_config = rbln_compile_configs[0]

optimum/rbln/transformers/models/grounding_dino/modeling_grounding_dino.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from pathlib import Path
-from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor, nn
@@ -206,7 +206,8 @@ class RBLNGroundingDinoForObjectDetection(RBLNModel):
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
     @classmethod
-    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
+    def get_pytorch_model(cls, *args, **kwargs):
+        model = super().get_pytorch_model(*args, **kwargs)
         model.encoder = model.model.encoder
         model.decoder = model.model.decoder
         model.text_backbone = model.model.text_backbone
@@ -216,7 +217,7 @@ class RBLNGroundingDinoForObjectDetection(RBLNModel):
         return model
 
     @classmethod
-    def _wrap_model_if_needed(
+    def wrap_model_if_needed(
         cls, model: torch.nn.Module, rbln_config: RBLNGroundingDinoForObjectDetectionConfig
     ) -> torch.nn.Module:
         return model.model.text_projection
@@ -529,26 +530,9 @@ class RBLNGroundingDinoForObjectDetection(RBLNModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        labels: List[Dict[str, Union[torch.LongTensor, torch.FloatTensor]]] = None,
         **kwargs,
-    ) -> Union[GroundingDinoObjectDetectionOutput, Tuple]:
-        """
-        Forward pass for the RBLN-optimized GroundingDinoForObjectDetection model.
-
-        Args:
-            pixel_values (torch.Tensor of shape (batch_size, num_channels, image_size, image_size)): The tensors corresponding to the input images.
-            input_ids (torch.LongTensor of shape (batch_size, text_sequence_length)): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it.
-            token_type_ids (torch.LongTensor of shape (batch_size, text_sequence_length), optional): Segment token indices to indicate first and second portions of the inputs.
-            attention_mask (torch.Tensor of shape (batch_size, sequence_length), optional): Mask to avoid performing attention on padding token indices.
-            pixel_mask (torch.Tensor of shape (batch_size, height, width), optional): Mask to avoid performing attention on padding pixel values.
-            encoder_outputs (Tuple consists of last_hidden_state of shape(batch_size, sequence_length, hidden_size), optional): A sequence of hidden-states at the output of the last layer of the encoder.
-            output_attentions (bool, optional): Whether or not to return the attentions tensors of all attention layers.
-            output_hidden_states (bool, optional): Whether or not to return the hidden states of all layers.
-            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 GroundingDinoObjectDetectionOutput object.
-        """
-
+    ):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         # Pad image to rbln_config.image_height and rbln_config.image_width
@@ -679,7 +663,7 @@ class RBLNGroundingDinoEncoder(RBLNModel):
         self.encoder_runtime = RBLNPytorchRuntime(self.model[0])
 
     @classmethod
-    def _wrap_model_if_needed(
+    def wrap_model_if_needed(
         cls, model: torch.nn.Module, rbln_config: RBLNGroundingDinoForObjectDetectionConfig
     ) -> torch.nn.Module:
         model = _GroundingDinoEncoder(model, rbln_config).eval()
@@ -877,7 +861,7 @@ class RBLNGroundingDinoDecoder(RBLNModel):
         self.decoder_runtime = RBLNPytorchRuntime(self.model[0])
 
     @classmethod
-    def _wrap_model_if_needed(
+    def wrap_model_if_needed(
         cls, model: torch.nn.Module, rbln_config: RBLNGroundingDinoForObjectDetectionConfig
     ) -> torch.nn.Module:
         return _GroundingDinoDecoder(model, rbln_config).eval()

optimum/rbln/transformers/models/idefics3/modeling_idefics3.py

@@ -110,7 +110,7 @@ class RBLNIdefics3VisionTransformer(RBLNModel):
         return self.embeddings
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
         class Idefics3VisionTransformerWrapper(torch.nn.Module):
             def __init__(self, model: "Idefics3VisionTransformer"):
                 super().__init__()
@@ -240,7 +240,9 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationM
         return True
 
     @classmethod
-    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
+    def get_pytorch_model(cls, *args, **kwargs):
+        model = super().get_pytorch_model(*args, **kwargs)
+
         with no_init_weights():
             model_cls_name = model.model.text_model.__class__.__name__
             causal_model_cls_name = model_cls_name.replace("Model", "ForCausalLM")
@@ -269,7 +271,7 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationM
         return self.text_model.get_input_embeddings()
 
     @classmethod
-    def _wrap_model_if_needed(cls, model, rbln_config):
+    def wrap_model_if_needed(cls, model, rbln_config):
         return model.model.connector
 
     @classmethod

optimum/rbln/transformers/models/llava/modeling_llava.py

@@ -88,22 +88,15 @@ class LoopVisionTower(LoopProcessor):
 
 
 class LoopProjector(LoopProcessor):
-    def __init__(self, multi_modal_projector: "RBLNModel", rbln_config=None):
+    def __init__(self, multi_modal_projector: "RBLNModel"):
         super().__init__(model=multi_modal_projector)
-        self.rbln_config = rbln_config
 
     def _get_batch_size(self, image_feature, **kwargs):
         return image_feature.shape[0]
 
     def _prepare_inputs_for_iteration(self, index, common_inputs, image_feature, **kwargs):
         image_feature_item = image_feature[index : index + 1]
-        if hasattr(self.rbln_config.vision_tower, "max_image_size"):
-            out_buffer = [
-                tensor[:, index * image_feature.shape[1] : (index + 1) * image_feature.shape[1], :]
-                for tensor in kwargs["out"]
-            ]
-        else:
-            out_buffer = [tensor[index : index + 1] for tensor in kwargs["out"]]
+        out_buffer = [tensor[index : index + 1] for tensor in kwargs["out"]]
         return ([image_feature_item], {"out": out_buffer})
 
     def _process_outputs(self, outputs: list, **kwargs):
@@ -182,7 +175,9 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
         return True
 
     @classmethod
-    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
+    def get_pytorch_model(cls, *args, **kwargs):
+        model = super().get_pytorch_model(*args, **kwargs)
+
         with no_init_weights():
             model_cls_name = model.model.language_model.__class__.__name__
             causal_model_cls_name = model_cls_name.replace("Model", "ForCausalLM")
@@ -199,7 +194,7 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
     def __post_init__(self, **kwargs):
         self.vision_tower = LoopVisionTower(self.rbln_submodules[0])
         self.language_model = self.rbln_submodules[1]
-        self.multi_modal_projector = LoopProjector(self.model[0], rbln_config=self.rbln_config)
+        self.multi_modal_projector = LoopProjector(self.model[0])
         self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
         return super().__post_init__(**kwargs)
 
@@ -213,7 +208,7 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
         return self.language_model.get_input_embeddings()
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
+    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         return model.multi_modal_projector
 
     @classmethod
@@ -226,8 +221,10 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
     ) -> RBLNModelConfig:
         # support for pixtral that needs padding
         if hasattr(rbln_config.vision_tower, "max_image_size"):
-            num_positions = (rbln_config.vision_tower.max_image_size[0] // model_config.vision_config.patch_size) * (
-                rbln_config.vision_tower.max_image_size[1] // model_config.vision_config.patch_size
+            num_positions = (
+                rbln_config.batch_size
+                * (rbln_config.vision_tower.max_image_size[0] // model_config.vision_config.patch_size)
+                * (rbln_config.vision_tower.max_image_size[1] // model_config.vision_config.patch_size)
             )
             selected_image_feature_dim = num_positions
 
@@ -356,32 +353,23 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
 
         if hasattr(self.rbln_config.vision_tower, "max_image_size"):
             num_real_patches = selected_image_feature.shape[1]
-            max_patches = (self.rbln_config.vision_tower.max_image_size[0] // self.config.vision_config.patch_size) * (
-                self.rbln_config.vision_tower.max_image_size[1] // self.config.vision_config.patch_size
+            max_patches = (
+                (self.rbln_config.vision_tower.max_image_size[0] // self.config.vision_config.patch_size)
+                * (self.rbln_config.vision_tower.max_image_size[1] // self.config.vision_config.patch_size)
+                * pixel_values.shape[0]
             )
+            num_padding_patches = max_patches - num_real_patches
 
-            chunks = []
-            for i in range(0, num_real_patches, max_patches):
-                chunk = selected_image_feature[:, i : i + max_patches, :]
-                chunk_size = chunk.shape[1]
-
-                if chunk_size < max_patches:
-                    padding_tensor = torch.zeros(
-                        (selected_image_feature.shape[0], max_patches - chunk_size, selected_image_feature.shape[2]),
-                        dtype=selected_image_feature.dtype,
-                    )
-                    chunk = torch.cat([chunk, padding_tensor], dim=1)
-                chunks.append(chunk)
-
-            split_features = torch.cat(chunks, dim=0)
-            num_chunks = len(chunks)
-            projector_out_size = [1, max_patches * num_chunks, self.config.text_config.hidden_size]
+            projector_out_size = [1, max_patches, self.config.text_config.hidden_size]
             projector_out_buffer = [torch.empty(size=projector_out_size, dtype=torch.float32, device="cpu")]
-            projected_features = self.multi_modal_projector(split_features, out=projector_out_buffer)
-            projected_features = projected_features.view(
-                selected_image_feature.shape[0], num_chunks * max_patches, self.config.text_config.hidden_size
+
+            padding_tensor = torch.zeros(
+                (selected_image_feature.shape[0], num_padding_patches, selected_image_feature.shape[2]),
+                dtype=selected_image_feature.dtype,
             )
-            image_features = projected_features[:, :num_real_patches, :]
+            padded_feature = torch.cat([selected_image_feature, padding_tensor], dim=1)
+            padded_projected_feature = self.multi_modal_projector(padded_feature, out=projector_out_buffer)
+            image_features = padded_projected_feature[:, :num_real_patches, :]
         else:
             projector_out_size = [
                 pixel_values.shape[0] * pixel_values.shape[1],

optimum/rbln/transformers/models/llava_next/modeling_llava_next.py

@@ -139,7 +139,9 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGeneration
         return True
 
     @classmethod
-    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
+    def get_pytorch_model(cls, *args, **kwargs):
+        model = super().get_pytorch_model(*args, **kwargs)
+
         with no_init_weights():
             model_cls_name = model.model.language_model.__class__.__name__
             causal_model_cls_name = model_cls_name.replace("Model", "ForCausalLM")
@@ -190,7 +192,7 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGeneration
         return self.language_model.get_input_embeddings()
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
+    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         return model.multi_modal_projector
 
     @classmethod

optimum/rbln/transformers/models/opt/modeling_opt.py

@@ -69,7 +69,7 @@ class RBLNOPTForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         return layer
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
+    def wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
         for i in range(len(model.model.decoder.layers)):
             model.model.decoder.layers[i] = cls.modify_opt_decoder_layer(model.model.decoder.layers[i])
 
@@ -95,7 +95,7 @@ class RBLNOPTModel(RBLNDecoderOnlyModel):
         return layer
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
+    def wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
         for i in range(len(model.decoder.layers)):
             model.decoder.layers[i] = cls.modify_opt_decoder_layer(model.decoder.layers[i])
 

optimum/rbln/transformers/models/pegasus/modeling_pegasus.py

@@ -54,7 +54,7 @@ class RBLNPegasusForConditionalGeneration(RBLNModelForSeq2SeqLM):
     support_causal_attn = True
 
     @classmethod
-    def _wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNPegasusForConditionalGenerationConfig):
+    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNPegasusForConditionalGenerationConfig):
         return PegasusWrapper(
             model, enc_max_seq_len=rbln_config.enc_max_seq_len, use_attention_mask=rbln_config.use_attention_mask
         )

optimum/rbln/transformers/models/pixtral/modeling_pixtral.py

@@ -229,7 +229,7 @@ class RBLNPixtralVisionModel(RBLNModel):
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
     @classmethod
-    def _wrap_model_if_needed(
+    def wrap_model_if_needed(
         cls, model: torch.nn.Module, rbln_config: RBLNPixtralVisionModelConfig
     ) -> torch.nn.Module:
         wrapper_cfg = {
@@ -293,18 +293,6 @@ class RBLNPixtralVisionModel(RBLNModel):
         return_dict: bool = True,
         **kwargs,
     ) -> Union[Tuple, BaseModelOutput]:
-        """
-        Forward pass for the RBLN-optimized Pixtral vision model.
-
-        Args:
-            pixel_values (torch.Tensor of shape (batch_size, num_channels, image_size, image_size)) — The tensors corresponding to the input images. Pixel values can be obtained using PixtralImageProcessor. See PixtralImageProcessor.call() for details (PixtralProcessor uses PixtralImageProcessor for processing images).
-            image_sizes (torch.Tensor of shape (batch_size, 2), optional) — The sizes of the images in the batch, being (height, width) for each image.
-            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.
-            return_dict (bool, optional) — Whether or not to return a ModelOutput instead of a plain tuple.
-
-        Returns:
-            BaseModelOutput or tuple(torch.FloatTensor)
-        """
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
         )

optimum/rbln/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py

@@ -88,7 +88,7 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
     @classmethod
-    def _wrap_model_if_needed(
+    def wrap_model_if_needed(
         cls, model: "PreTrainedModel", rbln_config: RBLNQwen2_5_VisionTransformerPretrainedModelConfig
     ):
         return Qwen2_5_VisionTransformerWrapper(model).eval()
@@ -393,7 +393,8 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         return True
 
     @classmethod
-    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
+    def get_pytorch_model(cls, *args, **kwargs):
+        model = super().get_pytorch_model(*args, **kwargs)
         model.model.lm_head = model.lm_head
         model.lm_head = None
         del model.lm_head

optimum/rbln/transformers/models/qwen2_vl/modeling_qwen2_vl.py

@@ -89,7 +89,7 @@ class RBLNQwen2VisionTransformerPretrainedModel(RBLNModel):
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
     @classmethod
-    def _wrap_model_if_needed(
+    def wrap_model_if_needed(
         cls, model: "PreTrainedModel", rbln_config: RBLNQwen2VisionTransformerPretrainedModelConfig
     ):
         return Qwen2VisionTransformerWrapper(model).eval()
@@ -282,7 +282,8 @@ class RBLNQwen2VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         return True
 
     @classmethod
-    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
+    def get_pytorch_model(cls, *args, **kwargs):
+        model = super().get_pytorch_model(*args, **kwargs)
         model.model.lm_head = model.lm_head
         model.lm_head = None
         del model.lm_head

optimum/rbln/transformers/models/resnet/configuration_resnet.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 
 
-from typing import Optional
-
 from ...configuration_generic import RBLNModelForImageClassificationConfig
 
 
@@ -25,18 +23,3 @@ class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConf
     This configuration class stores the configuration parameters specific to
     RBLN-optimized ResNet models for image classification tasks.
     """
-
-    def __init__(self, output_hidden_states: Optional[bool] = None, **kwargs):
-        """
-        Args:
-            image_size (Optional[Union[int, Tuple[int, int]]]): The size of input images.
-                Can be an integer for square images or a tuple (height, width).
-            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
-            output_hidden_states (bool, optional) — Whether or not to return the hidden states of all layers.
-            kwargs: Additional arguments passed to the parent RBLNModelConfig.
-
-        Raises:
-            ValueError: If batch_size is not a positive integer.
-        """
-        super().__init__(**kwargs)
-        self.output_hidden_states = output_hidden_states

optimum/rbln/transformers/models/resnet/modeling_resnet.py

@@ -13,17 +13,7 @@
 # limitations under the License.
 
 
-from typing import TYPE_CHECKING, Optional, Tuple, Union
-
-import torch
-from transformers.modeling_outputs import ImageClassifierOutputWithNoAttention
-
 from ...modeling_generic import RBLNModelForImageClassification
-from .configuration_resnet import RBLNResNetForImageClassificationConfig
-
-
-if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig, PreTrainedModel
 
 
 class RBLNResNetForImageClassification(RBLNModelForImageClassification):
@@ -34,66 +24,3 @@ class RBLNResNetForImageClassification(RBLNModelForImageClassification):
     on RBLN devices, supporting image classification with convolutional neural networks
     designed for computer vision tasks.
     """
-
-    @classmethod
-    def _update_rbln_config(
-        cls,
-        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
-        model: Optional["PreTrainedModel"] = None,
-        model_config: Optional["PretrainedConfig"] = None,
-        rbln_config: Optional["RBLNResNetForImageClassificationConfig"] = None,
-    ) -> "RBLNResNetForImageClassificationConfig":
-        if rbln_config.output_hidden_states is None:
-            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
-
-        rbln_config = super()._update_rbln_config(
-            preprocessors=preprocessors,
-            model=model,
-            model_config=model_config,
-            rbln_config=rbln_config,
-        )
-
-        return rbln_config
-
-    @classmethod
-    def _wrap_model_if_needed(
-        cls, model: torch.nn.Module, rbln_config: "RBLNResNetForImageClassificationConfig"
-    ) -> torch.nn.Module:
-        class _ResNetForImageClassification(torch.nn.Module):
-            def __init__(self, model: torch.nn.Module, output_hidden_states: bool):
-                super().__init__()
-                self.model = model
-                self.output_hidden_states = output_hidden_states
-
-            def forward(self, *args, **kwargs):
-                output = self.model(*args, output_hidden_states=self.output_hidden_states, **kwargs)
-                return output
-
-        return _ResNetForImageClassification(model, rbln_config.output_hidden_states)
-
-    def forward(
-        self, pixel_values: torch.Tensor, output_hidden_states: bool = None, return_dict: bool = None, **kwargs
-    ) -> Union[Tuple, ImageClassifierOutputWithNoAttention]:
-        """
-        Foward pass for the RBLN-optimized ResNet model for image classification.
-
-        Args:
-            pixel_values (torch.FloatTensor of shape (batch_size, channels, height, width)): The tensors corresponding to the input images.
-            output_hidden_states (bool, *optional*, defaults to False): Whether or not to return the hidden states of all layers.
-                See hidden_states under returned tensors for more details.
-            return_dict (bool, *optional*, defaults to True): Whether to return a dictionary of outputs.
-
-        Returns:
-            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a ImageClassifierOutputWithNoAttention object.
-        """
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
-        )
-
-        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."
-            )
-
-        return super().forward(pixel_values=pixel_values, return_dict=return_dict, **kwargs)

optimum/rbln/transformers/models/roberta/modeling_roberta.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 MaskedLMOutput, SequenceClassifierOutput
-
 from ...modeling_generic import RBLNModelForMaskedLM, RBLNModelForSequenceClassification
 
 
@@ -31,19 +26,6 @@ class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
 
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
-    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> Union[Tuple, MaskedLMOutput]:
-        """
-        Forward pass for the RBLN-optimized RoBERTa model for masked language modeling tasks.
-
-        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 MaskedLMOutput object.
-        """
-        return super().forward(input_ids, attention_mask, **kwargs)
-
 
 class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     """
@@ -55,18 +37,3 @@ class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     """
 
     rbln_model_input_names = ["input_ids", "attention_mask"]
-
-    def forward(
-        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs
-    ) -> Union[Tuple, SequenceClassifierOutput]:
-        """
-        Forward pass for the RBLN-optimized RoBERTa model for sequence classification tasks.
-
-        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/models/seq2seq/configuration_seq2seq.py

@@ -15,7 +15,6 @@
 from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
-from ....utils.deprecation import deprecate_kwarg
 from ....utils.logging import get_logger
 
 
@@ -25,13 +24,13 @@ logger = get_logger()
 class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
     support_paged_attention = None
 
-    @deprecate_kwarg(old_name="pad_token_id", version="0.10.0")
     def __init__(
         self,
         batch_size: Optional[int] = None,
         enc_max_seq_len: Optional[int] = None,
         dec_max_seq_len: Optional[int] = None,
         use_attention_mask: Optional[bool] = None,
+        pad_token_id: Optional[int] = None,
         kvcache_num_blocks: Optional[int] = None,
         kvcache_block_size: Optional[int] = None,
         **kwargs: Any,
@@ -42,6 +41,7 @@ class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
             enc_max_seq_len (Optional[int]): Maximum sequence length for the encoder.
             dec_max_seq_len (Optional[int]): Maximum sequence length for the decoder.
             use_attention_mask (Optional[bool]): Whether to use attention masks during inference.
+            pad_token_id (Optional[int]): The ID of the padding token in the vocabulary.
             kvcache_num_blocks (Optional[int]): The total number of blocks to allocate for the
                 PagedAttention KV cache for the SelfAttention. Defaults to batch_size.
             kvcache_block_size (Optional[int]): Sets the size (in number of tokens) of each block
@@ -61,6 +61,8 @@ class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
 
         self.use_attention_mask = use_attention_mask
 
+        self.pad_token_id = pad_token_id
+
         if self.support_paged_attention:
             self.kvcache_num_blocks = kvcache_num_blocks
             self.kvcache_block_size = kvcache_block_size