optimum-rbln 0.8.2a4__py3-none-any.whl → 0.8.2a5__py3-none-any.whl

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_gemma import RBLNGemmaForCausalLMConfig
-from .modeling_gemma import RBLNGemmaForCausalLM
+from .configuration_gemma import RBLNGemmaForCausalLMConfig, RBLNGemmaModelConfig
+from .modeling_gemma import RBLNGemmaForCausalLM, RBLNGemmaModel

optimum/rbln/transformers/models/gemma/configuration_gemma.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNGemmaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -40,3 +40,11 @@ class RBLNGemmaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     )
     ```
     """
+
+
+class RBLNGemmaModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN Gemma models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+    """

optimum/rbln/transformers/models/gemma/modeling_gemma.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from .gemma_architecture import GemmaWrapper
 
 
@@ -81,3 +81,15 @@ class RBLNGemmaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
 
     _decoder_wrapper_cls = GemmaWrapper
+
+
+class RBLNGemmaModel(RBLNDecoderOnlyModel):
+    """
+    The Gemma Model transformer without a language modeling head.
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based GemmaModel model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers GemmaModel model into a RBLN transformer model by:
+    """
+
+    _decoder_wrapper_cls = GemmaWrapper

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

@@ -13,8 +13,6 @@
 # limitations under the License.
 from typing import Any, Dict, Optional
 
-import rebel
-
 from ....configuration_utils import RBLNModelConfig
 from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
 from ..siglip.configuration_siglip import RBLNSiglipVisionModelConfig
@@ -39,10 +37,6 @@ class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
         )
         self.image_prefill_chunk_size = image_prefill_chunk_size
 
-        npu = self.npu or rebel.get_npu_name()
-        if npu == "RBLN-CA02":
-            raise NotImplementedError("Gemma3 is currently not supported on RBLN-CA02")
-
 
 class RBLNGemma3ForConditionalGenerationConfig(RBLNModelConfig):
     submodules = ["vision_tower", "language_model"]

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

@@ -31,7 +31,11 @@ from transformers.models.gemma3.modeling_gemma3 import Gemma3TextScaledWordEmbed
 
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
-from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyOutput, RBLNRuntimeModel
+from ..decoderonly.modeling_decoderonly import (
+    RBLNDecoderOnlyForCausalLMOutput,
+    RBLNDecoderOnlyModelForCausalLM,
+    RBLNRuntimeModel,
+)
 from .configuration_gemma3 import RBLNGemma3ForCausalLMConfig
 from .gemma3_architecture import Gemma3ForCausalLMWrapper
 
@@ -41,7 +45,7 @@ if TYPE_CHECKING:
 
 
 @dataclass
-class RBLNGemma3ForCausalLMOutput(RBLNDecoderOnlyOutput):
+class RBLNGemma3ForCausalLMOutput(RBLNDecoderOnlyForCausalLMOutput):
     attention_mask: Optional[torch.Tensor] = None
 
 
@@ -197,7 +201,7 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
 
     def _update_model_kwargs_for_generation(
         self,
-        outputs: RBLNDecoderOnlyOutput,
+        outputs: RBLNDecoderOnlyForCausalLMOutput,
         model_kwargs: Dict[str, Any],
         **kwargs,
     ) -> Dict[str, Any]:
@@ -266,7 +270,7 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
         position_ids: Optional[torch.Tensor] = None,
         token_type_ids: Optional[torch.Tensor] = None,
         **lm_kwargs: Dict[str, Any],
-    ) -> Union[Tuple, RBLNDecoderOnlyOutput]:
+    ) -> Union[Tuple, RBLNDecoderOnlyForCausalLMOutput]:
         # prefill
         if cache_position is None:
             logits = []
@@ -304,7 +308,7 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
                 position_ids=position_ids if self.rbln_config.language_model.use_position_ids else None,
             ).logits
 
-        return RBLNDecoderOnlyOutput(
+        return RBLNDecoderOnlyForCausalLMOutput(
             logits=logits, generate_idx=generate_idx, padded_cache_lengths=padded_cache_lengths
         )
 
@@ -509,7 +513,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
 
         logits = self.decode(inputs, cache_position, block_tables, local_block_tables, attention_mask, position_ids)
 
-        return RBLNDecoderOnlyOutput(logits=logits)
+        return RBLNDecoderOnlyForCausalLMOutput(logits=logits)
 
 
 class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_gpt2 import RBLNGPT2LMHeadModelConfig
-from .modeling_gpt2 import RBLNGPT2LMHeadModel
+from .configuration_gpt2 import RBLNGPT2LMHeadModelConfig, RBLNGPT2ModelConfig
+from .modeling_gpt2 import RBLNGPT2LMHeadModel, RBLNGPT2Model

optimum/rbln/transformers/models/gpt2/configuration_gpt2.py

@@ -12,11 +12,39 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNGPT2LMHeadModelConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     """
-    Configuration class for GPT-2 causal language model.
-    Inherits from RBLNDecoderOnlyModelForCausalLMConfig with no additional parameters.
+    Configuration class for RBLN GPT2 models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+    """
+
+
+class RBLNGPT2ModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN GPT2 models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNGPT2Model, RBLNGPT2ModelConfig
+
+    # Create a configuration object
+    config = RBLNGPT2ModelConfig(
+        batch_size=1,
+        max_seq_len=1024,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNGPT2Model.from_pretrained(
+        "openai/gpt2",
+        export=True,
+        rbln_config=config
+    )
+    ```
     """

optimum/rbln/transformers/models/gpt2/gpt2_architecture.py

@@ -13,11 +13,10 @@
 # limitations under the License.
 
 import math
-from typing import TYPE_CHECKING, Tuple
+from typing import TYPE_CHECKING, Tuple, Union
 
 import torch
 import torch.nn as nn
-from transformers import PreTrainedModel
 
 from ..decoderonly.decoderonly_architecture import (
     DecoderOnlyAttention,
@@ -28,7 +27,7 @@ from ..decoderonly.decoderonly_architecture import (
 
 
 if TYPE_CHECKING:
-    from transformers import GPT2LMHeadModel
+    from transformers import GPT2LMHeadModel, GPT2Model
 
 
 class GPT2Wrapper(DecoderOnlyWrapper):
@@ -44,11 +43,11 @@ class GPT2Wrapper(DecoderOnlyWrapper):
     def get_attn_layer(self, layer: nn.Module):
         return layer.attn
 
-    def get_model_layer(self, causal_lm: "GPT2LMHeadModel"):
-        return causal_lm.transformer
+    def get_model_layer(self, model: Union["GPT2LMHeadModel", "GPT2Model"]):
+        return model.transformer if self.is_causal_lm else model
 
-    def get_decoder_layers(self, causal_lm: PreTrainedModel):
-        return causal_lm.transformer.h
+    def get_decoder_layers(self, model: Union["GPT2LMHeadModel", "GPT2Model"]):
+        return model.transformer.h if self.is_causal_lm else model.h
 
 
 class GPT2Model(DecoderOnlyModel):

optimum/rbln/transformers/models/gpt2/modeling_gpt2.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from .gpt2_architecture import GPT2Wrapper
 
 
@@ -36,3 +36,18 @@ class RBLNGPT2LMHeadModel(RBLNDecoderOnlyModelForCausalLM):
 
     _decoder_wrapper_cls = GPT2Wrapper
     _use_rotary_emb = False
+
+
+class RBLNGPT2Model(RBLNDecoderOnlyModel):
+    """
+    The GPT2 Model transformer without a language modeling head.
+
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model.
+
+    A class to convert and run pre-trained transformers based GPT2Model model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers GPT2Model model into a RBLN transformer model by:
+    """
+
+    _decoder_wrapper_cls = GPT2Wrapper
+    _use_rotary_emb = False

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

@@ -35,7 +35,7 @@ from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ....utils.runtime_utils import RBLNPytorchRuntime
 from ..decoderonly.modeling_decoderonly import (
-    RBLNDecoderOnlyOutput,
+    RBLNDecoderOnlyForCausalLMOutput,
 )
 
 
@@ -494,7 +494,7 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
         if not return_dict:
             return logits, generate_idx
         else:
-            return RBLNDecoderOnlyOutput(
+            return RBLNDecoderOnlyForCausalLMOutput(
                 logits=logits,
                 generate_idx=generate_idx,
             )

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_llama import RBLNLlamaForCausalLMConfig
-from .modeling_llama import RBLNLlamaForCausalLM
+from .configuration_llama import RBLNLlamaForCausalLMConfig, RBLNLlamaModelConfig
+from .modeling_llama import RBLNLlamaForCausalLM, RBLNLlamaModel

optimum/rbln/transformers/models/llama/configuration_llama.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNLlamaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -40,3 +40,11 @@ class RBLNLlamaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     )
     ```
     """
+
+
+class RBLNLlamaModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN Llama models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+    """

optimum/rbln/transformers/models/llama/modeling_llama.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from .llama_architecture import LlamaWrapper
 
 
@@ -81,3 +81,15 @@ class RBLNLlamaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
 
     _decoder_wrapper_cls = LlamaWrapper
+
+
+class RBLNLlamaModel(RBLNDecoderOnlyModel):
+    """
+    The Llama Model transformer with a language modeling head (linear layer) on top.
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based LlamaModel model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers LlamaModel model into a RBLN transformer model by:
+    """
+
+    _decoder_wrapper_cls = LlamaWrapper

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

@@ -29,7 +29,7 @@ from transformers.modeling_outputs import BaseModelOutputWithPooling
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ....utils.logging import get_logger
-from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyOutput
+from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyForCausalLMOutput
 
 
 logger = get_logger(__name__)
@@ -258,7 +258,7 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
 
     def _update_model_kwargs_for_generation(
         self,
-        outputs: RBLNDecoderOnlyOutput,
+        outputs: RBLNDecoderOnlyForCausalLMOutput,
         model_kwargs: Dict[str, Any],
         **kwargs,
     ) -> Dict[str, Any]:
@@ -359,7 +359,7 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
         generate_idx: Optional[torch.Tensor] = None,
         batch_idx: Optional[int] = None,
         **kwargs,
-    ) -> Union[Tuple, RBLNDecoderOnlyOutput]:
+    ) -> Union[Tuple, RBLNDecoderOnlyForCausalLMOutput]:
         vision_feature_layer = (
             vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
         )
@@ -418,7 +418,7 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
                 cache_position=cache_position,
             )
             logits = output.logits
-        return RBLNDecoderOnlyOutput(logits=logits, generate_idx=generate_idx)
+        return RBLNDecoderOnlyForCausalLMOutput(logits=logits, generate_idx=generate_idx)
 
     # Almost copied from : https://github.com/huggingface/transformers/blob/6b550462139655d488d4c663086a63e98713c6b9/src/transformers/models/llava_next/modeling_llava_next.py
     def pack_image_features(self, image_features, image_sizes, vision_feature_select_strategy, image_newline=None):

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_mistral import RBLNMistralForCausalLMConfig
-from .modeling_mistral import RBLNMistralForCausalLM
+from .configuration_mistral import RBLNMistralForCausalLMConfig, RBLNMistralModelConfig
+from .modeling_mistral import RBLNMistralForCausalLM, RBLNMistralModel

optimum/rbln/transformers/models/mistral/configuration_mistral.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNMistralForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -40,3 +40,11 @@ class RBLNMistralForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     )
     ```
     """
+
+
+class RBLNMistralModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN Mistral models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+    """

optimum/rbln/transformers/models/mistral/mistral_architecture.py

@@ -15,5 +15,5 @@
 from ..decoderonly.decoderonly_architecture import DecoderOnlyWrapper
 
 
-class MistralForCausalLMWrapper(DecoderOnlyWrapper):
+class MistralWrapper(DecoderOnlyWrapper):
     pass

optimum/rbln/transformers/models/mistral/modeling_mistral.py

@@ -15,8 +15,12 @@
 from transformers import PretrainedConfig
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyModelForCausalLMConfig
-from .mistral_architecture import MistralForCausalLMWrapper
+from ...models.decoderonly import (
+    RBLNDecoderOnlyModel,
+    RBLNDecoderOnlyModelForCausalLM,
+    RBLNDecoderOnlyModelForCausalLMConfig,
+)
+from .mistral_architecture import MistralWrapper
 
 
 logger = logging.get_logger(__name__)
@@ -79,7 +83,26 @@ class RBLNMistralForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         ```
     """
 
-    _decoder_wrapper_cls = MistralForCausalLMWrapper
+    _decoder_wrapper_cls = MistralWrapper
+
+    @classmethod
+    def _update_sliding_window_config(
+        cls, model_config: PretrainedConfig, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
+    ):
+        rbln_config.cache_impl = "sliding_window"
+        rbln_config.sliding_window = model_config.sliding_window
+        rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
+
+        return rbln_config
+
+
+class RBLNMistralModel(RBLNDecoderOnlyModel):
+    """
+    The Mistral Model transformer without a language modeling head.
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    """
+
+    _decoder_wrapper_cls = MistralWrapper
 
     @classmethod
     def _update_sliding_window_config(

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_opt import RBLNOPTForCausalLMConfig
-from .modeling_opt import RBLNOPTForCausalLM
+from .configuration_opt import RBLNOPTForCausalLMConfig, RBLNOPTModelConfig
+from .modeling_opt import RBLNOPTForCausalLM, RBLNOPTModel

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNOPTForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -20,3 +20,10 @@ class RBLNOPTForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     Configuration class for OPT causal language model.
     Inherits from RBLNDecoderOnlyModelForCausalLMConfig with no additional parameters.
     """
+
+
+class RBLNOPTModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for OPT model.
+    Inherits from RBLNDecoderOnlyModelConfig with no additional parameters.
+    """

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

@@ -16,7 +16,7 @@ import torch.nn as nn
 from transformers import PreTrainedModel
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from ...models.decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
 from .opt_architecture import OPTWrapper
 
@@ -88,3 +88,43 @@ class RBLNOPTForCausalLM(RBLNDecoderOnlyModelForCausalLM):
             model.model.decoder.layers[i] = cls.modify_opt_decoder_layer(model.model.decoder.layers[i])
 
         return cls._decoder_wrapper_cls(model, **wrapper_cfg).eval()
+
+
+class RBLNOPTModel(RBLNDecoderOnlyModel):
+    """
+    The OPT Model transformer without a language modeling head.
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    """
+
+    _decoder_wrapper_cls = OPTWrapper
+    _use_rotary_emb = False
+
+    def modify_opt_decoder_layer(layer):
+        mlp = MLP(layer.fc1, layer.fc2, layer.activation_fn)
+        layer.mlp = mlp
+        del layer.fc1
+        del layer.fc2
+        del layer.activation_fn
+
+        return layer
+
+    @classmethod
+    def wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
+        wrapper_cfg = {
+            "max_seq_len": rbln_config.max_seq_len,
+            "attn_impl": rbln_config.attn_impl,
+            "kvcache_partition_len": rbln_config.kvcache_partition_len,
+            "kvcache_block_size": rbln_config.kvcache_block_size,
+            "use_rotary_emb": cls._use_rotary_emb,
+            "use_attention_mask": rbln_config.use_attention_mask,
+            "use_position_ids": rbln_config.use_position_ids,
+            "use_inputs_embeds": rbln_config.use_inputs_embeds,
+            "cache_impl": rbln_config.cache_impl,
+            "sliding_window": rbln_config.sliding_window,
+            "sliding_window_layers": rbln_config.sliding_window_layers,
+        }
+
+        for i in range(len(model.decoder.layers)):
+            model.decoder.layers[i] = cls.modify_opt_decoder_layer(model.decoder.layers[i])
+
+        return cls._decoder_wrapper_cls(model, **wrapper_cfg).eval()

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

@@ -40,11 +40,11 @@ class OPTWrapper(DecoderOnlyWrapper):
     def get_rbln_model_class(self):
         return OPTModel
 
-    def get_model_layer(self, causal_lm: "OPTForCausalLM"):
-        return causal_lm.model.decoder
+    def get_model_layer(self, model: "OPTForCausalLM"):
+        return model.model.decoder if self.is_causal_lm else model.decoder
 
-    def get_decoder_layers(self, causal_lm: "OPTForCausalLM"):
-        return causal_lm.model.decoder.layers
+    def get_decoder_layers(self, model: "OPTForCausalLM"):
+        return model.model.decoder.layers if self.is_causal_lm else model.decoder.layers
 
 
 class OPTAttention(DecoderOnlyAttention):

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

@@ -0,0 +1,17 @@
+# 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 ....ops import paged_attn_decode, paged_causal_attn_decode
+from .configuration_pegasus import RBLNPegasusForConditionalGenerationConfig, RBLNPegasusModelConfig
+from .modeling_pegasus import RBLNPegasusForConditionalGeneration, RBLNPegasusModel

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

@@ -0,0 +1,34 @@
+# 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 RBLNTransformerEncoderForFeatureExtractionConfig
+from ..seq2seq import RBLNModelForSeq2SeqLMConfig
+
+
+class RBLNPegasusModelConfig(RBLNTransformerEncoderForFeatureExtractionConfig):
+    """
+    Configuration class for RBLNPegasusModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized PEGASUS models for feature extraction tasks.
+    """
+
+
+class RBLNPegasusForConditionalGenerationConfig(RBLNModelForSeq2SeqLMConfig):
+    """
+    Configuration class for RBLNPegasusForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized PEGASUS models for conditional text generation tasks.
+    """