optimum-rbln 0.7.4a4__py3-none-any.whl → 0.7.4a6__py3-none-any.whl

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

@@ -13,48 +13,21 @@
 # limitations under the License.
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable
 
 import torch
-from transformers import (
-    AutoModelForTextEncoding,
-    PretrainedConfig,
-    T5EncoderModel,
-    T5ForConditionalGeneration,
-)
-from transformers.modeling_outputs import BaseModelOutput
-
-from ....diffusers.modeling_diffusers import RBLNDiffusionMixin
-from ....modeling import RBLNModel
-from ....modeling_config import RBLNCompileConfig, RBLNConfig
-from ....utils.logging import get_logger
-from ....utils.runtime_utils import RBLNPytorchRuntime
+from transformers import AutoModelForTextEncoding, T5EncoderModel, T5ForConditionalGeneration
+
+from ...modeling_generic import RBLNTransformerEncoderForFeatureExtraction
 from ...models.seq2seq import RBLNModelForSeq2SeqLM
+from .configuration_t5 import RBLNT5EncoderModelConfig, RBLNT5ForConditionalGenerationConfig
 from .t5_architecture import T5Wrapper
 
 
-logger = get_logger()
-
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
-
-
-class RBLNRuntimeModel(RBLNPytorchRuntime):
-    def forward(
-        self,
-        input_ids: torch.LongTensor,
-        attention_mask: torch.FloatTensor,
-        head_mask: torch.FloatTensor,
-        inputs_embeds: torch.FloatTensor,
-        **kwargs,
-    ):
-        return super().forward(
-            input_ids,
-            attention_mask,
-            head_mask,
-            inputs_embeds,
-            **kwargs,
-        )
+    from transformers import PreTrainedModel
+
+    from ....diffusers.modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
 
 
 class T5EncoderWrapper(torch.nn.Module):
@@ -67,136 +40,35 @@ class T5EncoderWrapper(torch.nn.Module):
         return self.model(*args, **kwargs, return_dict=False)
 
 
-class RBLNT5EncoderModel(RBLNModel):
+class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
     auto_model_class = AutoModelForTextEncoding
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
-    def __post_init__(self, **kwargs):
-        self.model = RBLNRuntimeModel(runtime=self.model[0])
-
     @classmethod
-    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
+    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNT5EncoderModelConfig):
         return T5EncoderWrapper(model)
 
     @classmethod
-    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        batch_size = rbln_config.get("batch_size", 1)
-        max_sequence_length = rbln_config.get("max_sequence_length", 256)
-        model_input_names = ["input_ids"]
-
-        rbln_config.update(
-            {
-                "batch_size": batch_size,
-                "max_seq_len": max_sequence_length,
-                "model_input_names": model_input_names,
-            }
-        )
-
-        return rbln_config
-
-    @classmethod
-    def _get_rbln_config(
+    def update_rbln_config_using_pipe(
         cls,
-        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
-        model_config: Optional["PretrainedConfig"] = None,
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
-        rbln_model_input_names = rbln_kwargs.get("model_input_names", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-
-        max_position_embeddings = getattr(model_config, "n_positions", None)
-
-        if rbln_max_seq_len is None:
-            rbln_max_seq_len = max_position_embeddings
-            if rbln_max_seq_len is None:
-                for tokenizer in preprocessors:
-                    if hasattr(tokenizer, "model_max_length"):
-                        rbln_max_seq_len = tokenizer.model_max_length
-                        break
-                if rbln_max_seq_len is None:
-                    raise ValueError("`rbln_max_seq_len` should be specified!")
-
-        if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
-            raise ValueError("`rbln_max_seq_len` should be less or equal than max_position_embeddings!")
-
-        signature_params = inspect.signature(cls.get_hf_class().forward).parameters.keys()
-
-        if rbln_model_input_names is None:
-            for tokenizer in preprocessors:
-                if hasattr(tokenizer, "model_input_names"):
-                    rbln_model_input_names = [name for name in signature_params if name in tokenizer.model_input_names]
-
-                    invalid_params = set(rbln_model_input_names) - set(signature_params)
-                    if invalid_params:
-                        raise ValueError(f"Invalid model input names: {invalid_params}")
-                    break
-            if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
-                rbln_model_input_names = cls.rbln_model_input_names
-            elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
-                raise ValueError(
-                    "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
-                    f"and be sure to make the order of the inputs same as T5EncoderModel forward() arguments like ({list(signature_params)})"
-                )
-        else:
-            invalid_params = set(rbln_model_input_names) - set(signature_params)
-            if invalid_params:
-                raise ValueError(f"Invalid model input names: {invalid_params}")
-            rbln_model_input_names = [name for name in signature_params if name in rbln_model_input_names]
-
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        input_info = [
-            (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
-            for model_input_name in rbln_model_input_names
-        ]
-
-        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-
-        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        pipe: "RBLNDiffusionMixin",
+        rbln_config: "RBLNDiffusionMixinConfig",
+        submodule_name: str,
+    ) -> "RBLNDiffusionMixinConfig":
+        submodule_config = getattr(rbln_config, submodule_name)
+        submodule_config.max_seq_len = rbln_config.max_seq_len or 256
+        submodule_config.model_input_names = ["input_ids"]
         return rbln_config
 
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple[torch.FloatTensor], BaseModelOutput]:
-        encoder_outputs = self.model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            inputs_embeds=inputs_embeds,
-            head_mask=head_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-        if not return_dict:
-            return (encoder_outputs,)
-        else:
-            return BaseModelOutput(last_hidden_state=encoder_outputs)
-
 
 class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
     support_causal_attn = False
 
     @classmethod
-    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
-        enc_max_seq_len = rbln_config.model_cfg["enc_max_seq_len"]
-        dec_max_seq_len = rbln_config.model_cfg["dec_max_seq_len"]
-
-        return T5Wrapper(model, enc_max_seq_len=enc_max_seq_len, dec_max_seq_len=dec_max_seq_len)
+    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
+        )
 
     def __getattr__(self, __name: str) -> Any:
         def redirect(func):

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

@@ -22,4 +22,5 @@
 # from Rebellions Inc.
 
 from ....ops import paged_add_softmax_attn_decode, rbln_cache_update
+from .configuration_time_series_transformer import RBLNTimeSeriesTransformerForPredictionConfig
 from .modeling_time_series_transformers import RBLNTimeSeriesTransformerForPrediction

optimum/rbln/transformers/models/time_series_transformers/configuration_time_series_transformer.py

@@ -0,0 +1,34 @@
+from typing import Optional
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNTimeSeriesTransformerForPredictionConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        enc_max_seq_len: Optional[int] = None,
+        dec_max_seq_len: Optional[int] = None,
+        num_parallel_samples: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            enc_max_seq_len (Optional[int]): Maximum sequence length for the encoder.
+            dec_max_seq_len (Optional[int]): Maximum sequence length for the decoder.
+            num_parallel_samples (Optional[int]): Number of samples to generate in parallel during prediction.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+
+        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}")
+
+        self.enc_max_seq_len = enc_max_seq_len
+        self.dec_max_seq_len = dec_max_seq_len
+        self.num_parallel_samples = num_parallel_samples

optimum/rbln/transformers/models/time_series_transformers/modeling_time_series_transformers.py

@@ -25,7 +25,7 @@ import inspect
 import logging
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union
 
 import rebel
 import torch
@@ -38,9 +38,10 @@ from transformers import (
 from transformers.modeling_outputs import ModelOutput, SampleTSPredictionOutput, Seq2SeqTSModelOutput
 from transformers.modeling_utils import no_init_weights
 
+from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
-from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.runtime_utils import RBLNPytorchRuntime
+from .configuration_time_series_transformer import RBLNTimeSeriesTransformerForPredictionConfig
 from .time_series_transformers_architecture import TimeSeriesTransformersWrapper
 
 
@@ -124,9 +125,9 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
 
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
-        self.batch_size = self.rbln_config.model_cfg["batch_size"]
-        self.dec_max_seq_len = self.rbln_config.model_cfg["dec_max_seq_len"]
-        self.num_parallel_samples = self.rbln_config.model_cfg["num_parallel_samples"]
+        self.batch_size = self.rbln_config.batch_size
+        self.dec_max_seq_len = self.rbln_config.dec_max_seq_len
+        self.num_parallel_samples = self.rbln_config.num_parallel_samples
 
         with no_init_weights():
             self._origin_model = TimeSeriesTransformerForPrediction._from_config(self.config)
@@ -156,12 +157,14 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
             return redirect(val)
 
     @classmethod
-    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
-        return TimeSeriesTransformersWrapper(model, rbln_config.model_cfg["num_parallel_samples"])
+    def wrap_model_if_needed(
+        self, model: "PreTrainedModel", rbln_config: RBLNTimeSeriesTransformerForPredictionConfig
+    ):
+        return TimeSeriesTransformersWrapper(model, rbln_config.num_parallel_samples)
 
     @classmethod
     @torch.inference_mode()
-    def get_compiled_model(cls, model, rbln_config: RBLNConfig):
+    def get_compiled_model(cls, model, rbln_config: RBLNTimeSeriesTransformerForPredictionConfig):
         wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
 
         enc_compile_config = rbln_config.compile_cfgs[0]
@@ -206,7 +209,7 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
         model: "PreTrainedModel",
         save_dir_path: Path,
         subfolder: str,
-        rbln_config: RBLNConfig,
+        rbln_config: RBLNTimeSeriesTransformerForPredictionConfig,
     ):
         """
         If you are unavoidably running on a CPU rather than an RBLN device,
@@ -217,31 +220,28 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
-        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
-        model_config: "PretrainedConfig",
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_batch_size = rbln_kwargs.get("batch_size", 1)
-        rbln_dec_max_seq_len = rbln_kwargs.get("dec_max_seq_len", None)
-        rbln_num_parallel_samples = rbln_kwargs.get("num_parallel_samples", None)
-
-        if not isinstance(rbln_batch_size, int):
-            raise TypeError(f"Expected rbln_batch_size to be an int, but got {type(rbln_batch_size)}")
-
-        rbln_num_parallel_samples = (
-            model_config.num_parallel_samples if rbln_num_parallel_samples is None else rbln_num_parallel_samples
-        )
-        if rbln_dec_max_seq_len is None:
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
+        model: Optional["PreTrainedModel"] = None,
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_config: Optional[RBLNTimeSeriesTransformerForPredictionConfig] = None,
+    ) -> RBLNTimeSeriesTransformerForPredictionConfig:
+        rbln_config.num_parallel_samples = rbln_config.num_parallel_samples or model_config.num_parallel_samples
+
+        if rbln_config.dec_max_seq_len is None:
             predict_length = model_config.prediction_length
-            rbln_dec_max_seq_len = (
+            rbln_config.dec_max_seq_len = (
                 predict_length if predict_length % 64 == 0 else predict_length + (64 - predict_length % 64)
             )
 
         # model input info
         enc_input_info = [
-            ("inputs_embeds", [rbln_batch_size, model_config.context_length, model_config.feature_size], "float32"),
+            (
+                "inputs_embeds",
+                [rbln_config.batch_size, model_config.context_length, model_config.feature_size],
+                "float32",
+            ),
         ]
         enc_input_info.extend(
             [
@@ -249,7 +249,7 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
                     "cross_key_value_states",
                     [
                         model_config.decoder_layers * 2,
-                        rbln_batch_size,
+                        rbln_config.batch_size,
                         model_config.decoder_attention_heads,
                         model_config.context_length,
                         model_config.d_model // model_config.decoder_attention_heads,
@@ -260,8 +260,12 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
         )
 
         dec_input_info = [
-            ("inputs_embeds", [rbln_batch_size * rbln_num_parallel_samples, 1, model_config.feature_size], "float32"),
-            ("attention_mask", [1, rbln_dec_max_seq_len], "float32"),
+            (
+                "inputs_embeds",
+                [rbln_config.batch_size * rbln_config.num_parallel_samples, 1, model_config.feature_size],
+                "float32",
+            ),
+            ("attention_mask", [1, rbln_config.dec_max_seq_len], "float32"),
             ("cache_position", [], "int32"),
             ("block_tables", [1, 1], "int16"),
         ]
@@ -271,7 +275,7 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
                     "cross_key_value_states",
                     [
                         model_config.decoder_layers * 2,  # 4
-                        rbln_batch_size,  # 64
+                        rbln_config.batch_size,  # 64
                         model_config.decoder_attention_heads,  # 2
                         model_config.context_length,  # 24
                         model_config.d_model // model_config.decoder_attention_heads,  # 13
@@ -286,8 +290,10 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
                     f"self_key_value_states_{i}",
                     [
                         1,
-                        model_config.decoder_attention_heads * rbln_num_parallel_samples * rbln_batch_size,
-                        rbln_dec_max_seq_len,
+                        model_config.decoder_attention_heads
+                        * rbln_config.num_parallel_samples
+                        * rbln_config.batch_size,
+                        rbln_config.dec_max_seq_len,
                         model_config.d_model // model_config.encoder_attention_heads,
                     ],
                     "float32",
@@ -298,38 +304,30 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
         enc_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=enc_input_info)
         dec_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=dec_input_info)
 
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[enc_compile_config, dec_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-
-        rbln_config.model_cfg.update(
-            {
-                "batch_size": rbln_batch_size,
-                "num_parallel_samples": rbln_num_parallel_samples,
-                "dec_max_seq_len": rbln_dec_max_seq_len,
-            }
-        )
-
+        rbln_config.set_compile_cfgs([enc_compile_config, dec_compile_config])
         return rbln_config
 
     @classmethod
     def _create_runtimes(
         cls,
         compiled_models: List[rebel.RBLNCompiledModel],
-        rbln_device_map: Dict[str, int],
-        activate_profiler: Optional[bool] = None,
+        rbln_config: RBLNTimeSeriesTransformerForPredictionConfig,
     ) -> List[rebel.Runtime]:
-        if any(model_name not in rbln_device_map for model_name in ["encoder", "decoder"]):
+        if any(model_name not in rbln_config.device_map for model_name in ["encoder", "decoder"]):
             cls._raise_missing_compiled_file_error(["encoder", "decoder"])
 
         return [
-            compiled_models[0].create_runtime(
-                tensor_type="pt", device=rbln_device_map["encoder"], activate_profiler=activate_profiler
+            rebel.Runtime(
+                compiled_models[0],
+                tensor_type="pt",
+                device=rbln_config.device_map["encoder"],
+                activate_profiler=rbln_config.activate_profiler,
             ),
-            compiled_models[1].create_runtime(
-                tensor_type="pt", device=rbln_device_map["decoder"], activate_profiler=activate_profiler
+            rebel.Runtime(
+                compiled_models[1],
+                tensor_type="pt",
+                device=rbln_config.device_map["decoder"],
+                activate_profiler=rbln_config.activate_profiler,
             ),
         ]
 

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

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

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

@@ -0,0 +1,19 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...configuration_generic import RBLNModelForMaskedLMConfig
+
+
+class RBLNWav2Vec2ForCTCConfig(RBLNModelForMaskedLMConfig):
+    rbln_model_input_names = ["input_values"]

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

@@ -12,26 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Any, Dict, Union
 
 import torch
-from transformers import AutoModelForMaskedLM, PretrainedConfig, Wav2Vec2ForCTC
+from transformers import AutoModelForMaskedLM, Wav2Vec2ForCTC
 from transformers.modeling_outputs import CausalLMOutput
 
-from ....modeling import RBLNModel
-from ....modeling_config import RBLNCompileConfig, RBLNConfig
-from ....utils.logging import get_logger
-
-
-logger = get_logger(__name__)
-
-if TYPE_CHECKING:
-    from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        AutoTokenizer,
-        PretrainedConfig,
-    )
+from ...modeling_generic import RBLNModelForMaskedLM
+from .configuration_wav2vec import RBLNWav2Vec2ForCTCConfig
 
 
 class _Wav2Vec2(torch.nn.Module):
@@ -44,11 +31,11 @@ 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 [`RBLNModel`]. Check the superclass documentation for the generic methods the
+    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:
@@ -58,60 +45,10 @@ class RBLNWav2Vec2ForCTC(RBLNModel):
 
     main_input_name = "input_values"
     auto_model_class = AutoModelForMaskedLM
+    rbln_dtype = "float32"
+    output_class = CausalLMOutput
+    output_key = "logits"
 
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> 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 _get_rbln_config(
-        cls,
-        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
-        model_config: "PretrainedConfig",
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-
-        if rbln_max_seq_len is None:
-            for tokenizer in preprocessors:
-                if hasattr(tokenizer, "model_max_length"):
-                    rbln_max_seq_len = tokenizer.model_max_length
-                    break
-            if rbln_max_seq_len is None:
-                raise ValueError("`rbln_max_seq_len` should be specified!")
-
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        input_info = [
-            (
-                "input_values",
-                [
-                    rbln_batch_size,
-                    rbln_max_seq_len,
-                ],
-                "float32",
-            ),
-        ]
-
-        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-
-        rbln_config.model_cfg.update(
-            {
-                "max_seq_len": rbln_max_seq_len,
-                "batch_size": rbln_batch_size,
-            }
-        )
-
-        return rbln_config
-
-    def forward(self, input_values: "torch.Tensor", **kwargs):
-        outputs = super().forward(input_values, **kwargs)
-        return CausalLMOutput(logits=outputs)

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

@@ -13,4 +13,5 @@
 # limitations under the License.
 
 from ....ops import paged_add_softmax_attn_decode
+from .configuration_whisper import RBLNWhisperForConditionalGenerationConfig
 from .modeling_whisper import RBLNWhisperForConditionalGeneration

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

@@ -0,0 +1,64 @@
+# 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.
+
+import rebel
+
+from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger()
+
+
+class RBLNWhisperForConditionalGenerationConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: int = None,
+        token_timestamps: bool = None,
+        use_attention_mask: bool = None,
+        enc_max_seq_len: int = None,
+        dec_max_seq_len: int = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (int, optional): The batch size for inference. Defaults to 1.
+            token_timestamps (bool, optional): Whether to output token timestamps during generation. Defaults to False.
+            use_attention_mask (bool, optional): Whether to use attention masks during inference. This is automatically
+                set to True for RBLN-CA02 devices.
+            enc_max_seq_len (int, optional): Maximum sequence length for the encoder.
+            dec_max_seq_len (int, optional): Maximum sequence length for the decoder.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+
+        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}")
+
+        self.token_timestamps = token_timestamps or False
+        self.enc_max_seq_len = enc_max_seq_len
+        self.dec_max_seq_len = dec_max_seq_len
+
+        self.use_attention_mask = use_attention_mask
+        npu = self.npu or rebel.get_npu_name()
+        if npu == "RBLN-CA02":
+            if self.use_attention_mask is False:
+                logger.warning("Attention mask should be used with RBLN-CA02. Setting use_attention_mask to True.")
+            self.use_attention_mask = True
+        else:
+            self.use_attention_mask = self.use_attention_mask or False