optimum-rbln 0.8.0.post2__py3-none-any.whl → 0.8.1__py3-none-any.whl

optimum/rbln/transformers/models/distilbert/__init__.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_distilbert import RBLNDistilBertForQuestionAnsweringConfig
+from .modeling_distilbert import RBLNDistilBertForQuestionAnswering
+
+
+__all__ = ["RBLNDistilBertForQuestionAnsweringConfig", "RBLNDistilBertForQuestionAnswering"]

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

@@ -0,0 +1,24 @@
+# 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 RBLNModelForQuestionAnsweringConfig
+
+
+class RBLNDistilBertForQuestionAnsweringConfig(RBLNModelForQuestionAnsweringConfig):
+    """
+    Configuration class for RBLNDistilBertForQuestionAnswering.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized DistilBERT models for question answering tasks.
+    """

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

@@ -0,0 +1,27 @@
+# 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 ...modeling_generic import RBLNModelForQuestionAnswering
+
+
+class RBLNDistilBertForQuestionAnswering(RBLNModelForQuestionAnswering):
+    """
+    RBLN optimized DistilBERT model for question answering tasks.
+
+    This class provides hardware-accelerated inference for DistilBERT models
+    on RBLN devices, supporting extractive question answering tasks where
+    the model predicts start and end positions of answers in text.
+    """
+
+    rbln_model_input_names = ["input_ids", "attention_mask"]

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

@@ -16,4 +16,9 @@ from ...configuration_generic import RBLNModelForDepthEstimationConfig
 
 
 class RBLNDPTForDepthEstimationConfig(RBLNModelForDepthEstimationConfig):
-    pass
+    """
+    Configuration class for RBLNDPTForDepthEstimation.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized DPT (Dense Prediction Transformer) models for depth estimation tasks.
+    """

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

@@ -17,4 +17,9 @@ from ...modeling_generic import RBLNModelForDepthEstimation
 
 
 class RBLNDPTForDepthEstimation(RBLNModelForDepthEstimation):
-    pass
+    """
+    RBLN optimized DPT model for depth estimation tasks.
+
+    This class provides hardware-accelerated inference for DPT (Dense Prediction Transformer)
+    models on RBLN devices, supporting monocular depth estimation from single images.
+    """

optimum/rbln/transformers/models/exaone/configuration_exaone.py

@@ -16,4 +16,27 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 
 class RBLNExaoneForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
-    pass
+    """
+    Configuration class for RBLN Exaone models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNExaoneForCausalLM, RBLNExaoneForCausalLMConfig
+
+    # Create a configuration object
+    config = RBLNExaoneForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=4096,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNExaoneForCausalLM.from_pretrained(
+        "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

optimum/rbln/transformers/models/exaone/exaone_architecture.py

@@ -60,7 +60,11 @@ class ExaoneForCausalLMWrapper(DecoderOnlyWrapper):
             new_layer = ExaoneLayer(layer, new_self_attn)
             new_layers.append(new_layer)
         new_model = ExaoneModel(
-            causal_lm.transformer, new_layers, partition_len=self.kvcache_partition_len, max_seq_len=max_seq_len
+            causal_lm.transformer,
+            new_layers,
+            partition_len=self.kvcache_partition_len,
+            max_seq_len=max_seq_len,
+            sliding_window_layers=self.sliding_window_layers,
         )
         new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
         return new_causal_lm

optimum/rbln/transformers/models/exaone/modeling_exaone.py

@@ -13,7 +13,11 @@
 # limitations under the License.
 
 
+import inspect
+from typing import Any, Callable
+
 from transformers import AutoModelForCausalLM
+from transformers.generation.utils import GenerationMixin
 
 from ....utils import logging
 from ..decoderonly import RBLNDecoderOnlyModelForCausalLM
@@ -25,22 +29,79 @@ logger = logging.get_logger(__name__)
 
 class RBLNExaoneForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
-    The Exaone Model transformer with a language modeling head on top (linear layer with weights tied to the input
-    embeddings).
+    The Exaone Model transformer with a language modeling head (linear layer) on top.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
 
-    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. 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 ExaoneForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers ExaoneForCausalLM model into a RBLN transformer model by:
 
-    It implements the methods to convert a pre-trained transformers Exaone model into a RBLN transformer model by:
     - transferring the checkpoint weights of the original into an optimized RBLN graph,
     - compiling the resulting graph using the RBLN compiler.
 
+    **Configuration:**
+    This model uses [`RBLNExaoneForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNExaoneForCausalLMConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNExaoneForCausalLMConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNExaoneForCausalLM
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNExaoneForCausalLM.from_pretrained(
+            "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 4096,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNExaoneForCausalLM.from_pretrained(
+            "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+
+        # Using a RBLNExaoneForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNExaoneForCausalLMConfig
+
+        config = RBLNExaoneForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=4096,
+            tensor_parallel_size=4
+        )
+        model = RBLNExaoneForCausalLM.from_pretrained(
+            "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct",
+            export=True,
+            rbln_config=config
+        )
+        ```
     """
 
     _decoder_wrapper_cls = ExaoneForCausalLMWrapper
     _hf_class = AutoModelForCausalLM
+    _supports_cache_class = True
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
         kwargs.setdefault("trust_remote_code", True)
         return super().from_pretrained(*args, **kwargs)
+
+    def __getattr__(self, __name: str) -> Any:
+        def redirect(func):
+            return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
+
+        val = getattr(GenerationMixin, __name)
+
+        if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
+            return redirect(val)
+        return val

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

@@ -16,4 +16,27 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 
 class RBLNGemmaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
-    pass
+    """
+    Configuration class for RBLN Gemma models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNGemmaForCausalLM, RBLNGemmaForCausalLMConfig
+
+    # Create a configuration object
+    config = RBLNGemmaForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=4096,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNGemmaForCausalLM.from_pretrained(
+        "google/gemma-7b",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

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

@@ -52,7 +52,11 @@ class GemmaWrapper(DecoderOnlyWrapper):
             new_layer = DecoderOnlyLayer(layer, new_self_attn)
             new_layers.append(new_layer)
         new_model = GemmaModel(
-            causal_lm.model, new_layers, partition_len=self.kvcache_partition_len, max_seq_len=max_seq_len
+            causal_lm.model,
+            new_layers,
+            partition_len=self.kvcache_partition_len,
+            max_seq_len=max_seq_len,
+            sliding_window_layers=self.sliding_window_layers,
         )
         new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
         return new_causal_lm

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

@@ -27,8 +27,57 @@ class RBLNGemmaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
     A class to convert and run pre-trained transformers based GemmaForCausalLM model on RBLN devices.
     It implements the methods to convert a pre-trained transformers GemmaForCausalLM model into a RBLN transformer model by:
+
     - transferring the checkpoint weights of the original into an optimized RBLN graph,
     - compiling the resulting graph using the RBLN compiler.
+
+    **Configuration:**
+    This model uses [`RBLNGemmaForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNGemmaForCausalLMConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNGemmaForCausalLMConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNGemmaForCausalLM
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNGemmaForCausalLM.from_pretrained(
+            "google/gemma-7b",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 4096,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNGemmaForCausalLM.from_pretrained(
+            "google/gemma-7b",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+
+        # Using a RBLNGemmaForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNGemmaForCausalLMConfig
+
+        config = RBLNGemmaForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=4096,
+            tensor_parallel_size=4
+        )
+        model = RBLNGemmaForCausalLM.from_pretrained(
+            "google/gemma-7b",
+            export=True,
+            rbln_config=config
+        )
+        ```
     """
 
     _decoder_wrapper_cls = GemmaWrapper

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

@@ -11,7 +11,7 @@
 # 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 typing import Optional
+from typing import Any, Dict, Optional
 
 import rebel
 
@@ -26,7 +26,7 @@ class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
         prefill_chunk_size: Optional[int] = None,
         use_position_ids: Optional[bool] = None,
         use_attention_mask: Optional[bool] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         # use_attention_mask and use_position_ids are always True for Gemma3
         use_attention_mask = use_attention_mask or True
@@ -53,7 +53,7 @@ class RBLNGemma3ForConditionalGenerationConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         vision_tower: Optional[RBLNModelConfig] = None,
         language_model: Optional[RBLNModelConfig] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args: