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

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

@@ -37,6 +37,7 @@ from ....utils.logging import get_logger
 from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyOutput
 from .configuration_qwen2_5_vl import (
     RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
+    RBLNQwen2_5_VLForConditionalGenerationConfig,
 )
 from .qwen2_5_vl_architecture import Qwen2_5_VisionTransformerWrapper, Qwen2_5_VL_LanguageModelWrapper
 
@@ -338,6 +339,40 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
 
 
 class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
+    """
+    RBLNQwen2_5_VLForConditionalGeneration is a multi-modal model that integrates vision and language processing capabilities,
+    optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.
+
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    Important Note:
+        This model includes a Large Language Model (LLM). For optimal performance, it is highly recommended to use
+        tensor parallelism for the language model. This can be achieved by using the `rbln_config` parameter in the
+        `from_pretrained` method. Refer to the `from_pretrained` documentation and the RBLNQwen2_5_VLForConditionalGenerationConfig class for details.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNQwen2_5_VLForConditionalGeneration
+
+        model = RBLNQwen2_5_VLForConditionalGeneration.from_pretrained(
+            "Qwen/Qwen2.5-VL-7B-Instruct",
+            export=True,
+            rbln_config={
+                "visual": {
+                    "max_seq_lens": 6400,
+                    "device": 0,
+                },
+                "tensor_parallel_size": 8,
+                "kvcache_partition_len": 16_384,
+                "max_seq_len": 114_688,
+                "device": [0, 1, 2, 3, 4, 5, 6, 7],
+            },
+        )
+
+        model.save_pretrained("compiled-qwen2.5-vl-7b-instruct")
+        ```
+    """
+
     auto_model_class = AutoModelForVision2Seq
     _rbln_submodules = [
         {"name": "visual"},
@@ -369,33 +404,19 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         cls,
         batch_size: int,
         query_length: int,
-        use_inputs_embeds: bool,
-        use_attention_mask: bool,
-        use_position_ids: bool,
-        max_seq_len: int,
-        kvcache_block_size: int,
-        kvcache_num_blocks: int,
-        num_key_value_heads: int,
-        num_hidden_layers: int,
-        hidden_size: int,
-        head_dim: int,
+        rbln_config: RBLNQwen2_5_VLForConditionalGenerationConfig,
+        model_config: PretrainedConfig,
     ):
-        input_info = super().get_input_info(
-            batch_size,
-            query_length,
-            use_inputs_embeds,
-            use_attention_mask,
-            use_position_ids,
-            max_seq_len,
-            kvcache_block_size,
-            kvcache_num_blocks,
-            num_key_value_heads,
-            num_hidden_layers,
-            hidden_size,
-            head_dim,
-        )
+        input_info = super().get_input_info(batch_size, query_length, rbln_config, model_config)
         pos_idx = 3
-        input_info.insert(pos_idx, ("position_emb", [2, batch_size, 1, query_length, head_dim], "float32"))
+        input_info.insert(
+            pos_idx,
+            (
+                "position_emb",
+                [2, batch_size, 1, query_length, model_config.hidden_size // model_config.num_attention_heads],
+                "float32",
+            ),
+        )
 
         return input_info
 

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

@@ -162,7 +162,8 @@ class Qwen2_5_VL_LanguageModelWrapper(DecoderOnlyWrapper):
         input_ids = None if self.use_inputs_embeds else args.pop(0)
         inputs_embeds = args.pop(0) if self.use_inputs_embeds else None
         cache_position = args.pop(0)
-        block_tables = args.pop(0)
+        global_block_tables = args.pop(0)
+        local_block_tables = None
         position_embeds = args.pop(0)
         query_position = args.pop(0) if self.phase == "prefill" else None
         position_ids = None
@@ -188,7 +189,8 @@ class Qwen2_5_VL_LanguageModelWrapper(DecoderOnlyWrapper):
             input_ids,
             inputs_embeds,
             cache_position,
-            block_tables,
+            global_block_tables,
+            local_block_tables,
             query_position,
             attention_mask,
             position_ids,

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

@@ -0,0 +1,23 @@
+# 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_resnet import RBLNResNetForImageClassificationConfig
+from .modeling_resnet import RBLNResNetForImageClassification
+
+
+__all__ = [
+    "RBLNResNetForImageClassificationConfig",
+    "RBLNResNetForImageClassification",
+]

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

@@ -0,0 +1,20 @@
+# 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 RBLNModelForImageClassificationConfig
+
+
+class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConfig):
+    ""

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

@@ -0,0 +1,22 @@
+# 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 RBLNModelForImageClassification
+
+
+class RBLNResNetForImageClassification(RBLNModelForImageClassification):
+    """
+    ResNet model for image classification tasks on RBLN NPU.
+    """

optimum/rbln/transformers/models/roberta/__init__.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_roberta import RBLNRobertaForMaskedLMConfig, RBLNRobertaForSequenceClassificationConfig
+from .modeling_roberta import RBLNRobertaForMaskedLM, RBLNRobertaForSequenceClassification
+
+
+__all__ = [
+    "RBLNRobertaForMaskedLMConfig",
+    "RBLNRobertaForSequenceClassificationConfig",
+    "RBLNRobertaForMaskedLM",
+    "RBLNRobertaForSequenceClassification",
+]

optimum/rbln/transformers/{configuration_alias.py → models/roberta/configuration_roberta.py}

@@ -12,38 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_generic import (
-    RBLNModelForAudioClassificationConfig,
-    RBLNModelForImageClassificationConfig,
-    RBLNModelForMaskedLMConfig,
-    RBLNModelForQuestionAnsweringConfig,
-    RBLNModelForSequenceClassificationConfig,
-)
-
-
-class RBLNASTForAudioClassificationConfig(RBLNModelForAudioClassificationConfig):
-    pass
-
-
-class RBLNDistilBertForQuestionAnsweringConfig(RBLNModelForQuestionAnsweringConfig):
-    pass
-
-
-class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConfig):
-    pass
-
-
-class RBLNXLMRobertaForSequenceClassificationConfig(RBLNModelForSequenceClassificationConfig):
-    pass
-
-
-class RBLNRobertaForSequenceClassificationConfig(RBLNModelForSequenceClassificationConfig):
-    pass
+from ...configuration_generic import RBLNModelForMaskedLMConfig, RBLNModelForSequenceClassificationConfig
 
 
 class RBLNRobertaForMaskedLMConfig(RBLNModelForMaskedLMConfig):
-    pass
+    ""
 
 
-class RBLNViTForImageClassificationConfig(RBLNModelForImageClassificationConfig):
-    pass
+class RBLNRobertaForSequenceClassificationConfig(RBLNModelForSequenceClassificationConfig):
+    ""

optimum/rbln/transformers/{modeling_alias.py → models/roberta/modeling_roberta.py}

@@ -12,42 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..utils.logging import get_logger
-from .modeling_generic import (
-    RBLNModelForAudioClassification,
-    RBLNModelForImageClassification,
-    RBLNModelForMaskedLM,
-    RBLNModelForQuestionAnswering,
-    RBLNModelForSequenceClassification,
-)
+from ...modeling_generic import RBLNModelForMaskedLM, RBLNModelForSequenceClassification
 
 
-logger = get_logger()
-
-
-class RBLNASTForAudioClassification(RBLNModelForAudioClassification):
-    pass
-
-
-class RBLNDistilBertForQuestionAnswering(RBLNModelForQuestionAnswering):
-    rbln_model_input_names = ["input_ids", "attention_mask"]
-
-
-class RBLNResNetForImageClassification(RBLNModelForImageClassification):
-    pass
-
-
-class RBLNXLMRobertaForSequenceClassification(RBLNModelForSequenceClassification):
+class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
 class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     rbln_model_input_names = ["input_ids", "attention_mask"]
-
-
-class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
-    rbln_model_input_names = ["input_ids", "attention_mask"]
-
-
-class RBLNViTForImageClassification(RBLNModelForImageClassification):
-    pass

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_seq2seq2 import RBLNModelForSeq2SeqLMConfig
+from .configuration_seq2seq import RBLNModelForSeq2SeqLMConfig
 from .modeling_seq2seq import RBLNModelForSeq2SeqLM

optimum/rbln/transformers/models/seq2seq/{configuration_seq2seq2.py → configuration_seq2seq.py}

@@ -12,7 +12,7 @@
 # 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
 
@@ -31,7 +31,7 @@ class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
         dec_max_seq_len: Optional[int] = None,
         use_attention_mask: Optional[bool] = None,
         pad_token_id: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

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

@@ -26,7 +26,7 @@ from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
 from ....utils.logging import get_logger
 from ....utils.runtime_utils import RBLNPytorchRuntime
-from .configuration_seq2seq2 import RBLNModelForSeq2SeqLMConfig
+from .configuration_seq2seq import RBLNModelForSeq2SeqLMConfig
 
 
 logger = get_logger(__name__)

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

@@ -24,6 +24,7 @@ class RBLNSiglipVisionModelConfig(RBLNModelConfig):
         image_size: Optional[int] = None,
         interpolate_pos_encoding: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
         **kwargs,
     ):
         """
@@ -33,6 +34,7 @@ class RBLNSiglipVisionModelConfig(RBLNModelConfig):
                 a tuple/list (height, width), or a dictionary with 'height' and 'width' keys.
             interpolate_pos_encoding (Optional[bool]): Whether to interpolate the position encoding.
             output_hidden_states: (Optional[bool]): Whether to return hidden states.
+            output_attentions: (Optional[bool]): Whether to return attentions.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
@@ -46,6 +48,7 @@ class RBLNSiglipVisionModelConfig(RBLNModelConfig):
         self.image_size = image_size
         self.interpolate_pos_encoding = interpolate_pos_encoding or False
         self.output_hidden_states = output_hidden_states
+        self.output_attentions = output_attentions
 
     @property
     def image_width(self):

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

@@ -12,12 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 import torch
 from transformers import SiglipVisionConfig, SiglipVisionModel
 from transformers.modeling_outputs import BaseModelOutputWithPooling
-from transformers.models.siglip.modeling_siglip import SiglipVisionModelOutput
 
 from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
@@ -34,11 +33,18 @@ if TYPE_CHECKING:
 
 
 class _SiglipVisionModel(torch.nn.Module):
-    def __init__(self, model: SiglipVisionModel, interpolate_pos_encoding: bool, output_hidden_states: bool):
+    def __init__(
+        self,
+        model: SiglipVisionModel,
+        interpolate_pos_encoding: bool,
+        output_hidden_states: bool,
+        output_attentions: bool,
+    ):
         super().__init__()
         self.vision_model = model.vision_model
         self.interpolate_pos_encoding = interpolate_pos_encoding
         self.output_hidden_states = output_hidden_states
+        self.output_attentions = output_attentions
 
     def forward(self, inp):
         enc_out = self.vision_model(
@@ -46,6 +52,7 @@ class _SiglipVisionModel(torch.nn.Module):
             output_hidden_states=self.output_hidden_states,
             return_dict=False,
             interpolate_pos_encoding=self.interpolate_pos_encoding,
+            output_attentions=self.output_attentions,
         )
         return tuple(x for x in enc_out if x is not None)
 
@@ -56,6 +63,7 @@ class RBLNSiglipVisionModel(RBLNModel):
         wrapper_cfg = {
             "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
             "output_hidden_states": rbln_config.output_hidden_states,
+            "output_attentions": rbln_config.output_attentions,
         }
         return _SiglipVisionModel(model, **wrapper_cfg).eval()
 
@@ -81,8 +89,10 @@ class RBLNSiglipVisionModel(RBLNModel):
         if rbln_config.image_size is None:
             raise ValueError("`rbln_image_size` should be specified!")
 
+        if rbln_config.output_attentions is None:
+            rbln_config.output_attentions = getattr(model_config, "output_attentions", False)
         if rbln_config.output_hidden_states is None:
-            rbln_config.output_hidden_states = model_config.output_hidden_states
+            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
 
         rbln_compile_config = RBLNCompileConfig(
             input_info=[
@@ -104,43 +114,74 @@ class RBLNSiglipVisionModel(RBLNModel):
 
     def forward(
         self,
-        pixel_values: Optional[torch.FloatTensor] = None,
+        pixel_values: torch.Tensor,
         return_dict: bool = None,
+        output_attentions: bool = None,
+        output_hidden_states: bool = None,
         interpolate_pos_encoding: bool = False,
-        **kwargs,
-    ) -> Union[Tuple, SiglipVisionModelOutput]:
-        if len(kwargs) > 0 and any(kwargs.values()):
-            logger.warning(f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__}.")
+        **kwargs: Dict[str, Any],
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        if len(kwargs) > 0 and any(value is not None for value in kwargs.values()):
+            logger.warning(
+                f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__.__name__}."
+            )
+
+        output_attentions = output_attentions if output_attentions is not None else self.rbln_config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        if output_attentions != self.rbln_config.output_attentions:
+            raise ValueError(
+                f"Variable output_attentions {output_attentions} is not equal to rbln_config.output_attentions {self.rbln_config.output_attentions} "
+                f"Please compile again with the correct argument."
+            )
+
+        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."
+            )
 
         if interpolate_pos_encoding != self.rbln_config.interpolate_pos_encoding:
             raise ValueError(
-                f"Variable interpolate_pos_encoding {interpolate_pos_encoding} is not equal to rbln_config.interpolate_pos_encoding {self.rbln_config.interpolate_pos_encoding}"
+                f"Variable interpolate_pos_encoding {interpolate_pos_encoding} is not equal to rbln_config.interpolate_pos_encoding {self.rbln_config.interpolate_pos_encoding} "
                 f"Please compile again with the correct argument."
             )
+
         output = super().forward(pixel_values, return_dict=return_dict)
         return output
 
     def _prepare_output(self, output, return_dict):
-        """
-        Prepare model output based on return_dict flag.
-        This method can be overridden by subclasses to provide task-specific output handling.
-        """
+        # Prepare model output based on return_dict flag.
+        # This method can be overridden by subclasses to provide task-specific output handling.
+
         if not return_dict:
             return (output,) if not isinstance(output, (tuple, list)) else output
         else:
-            last_hidden_state = (
-                output[0]
-                if self.rbln_config.interpolate_pos_encoding or self.rbln_config.output_hidden_states
-                else output
-            )
-            pooler_output = output[1] if self.rbln_config.interpolate_pos_encoding else None
+            last_hidden_state = output.pop(0) if isinstance(output, (tuple, list)) else output
+            vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+            pooler_output = output.pop(0) if getattr(vision_config, "vision_use_head", True) else None
+
             if self.rbln_config.output_hidden_states:
-                hidden_states = (output[2:] if self.rbln_config.interpolate_pos_encoding else output[1:],)
+                hidden_states = ()
+                num_hidden_layers = vision_config.num_hidden_layers
+                for _ in range(num_hidden_layers + 1):
+                    hidden_states += (output.pop(0),)
             else:
                 hidden_states = None
 
+            if self.rbln_config.output_attentions:
+                attentions = ()
+                num_hidden_layers = vision_config.num_hidden_layers
+                for _ in range(num_hidden_layers):
+                    attentions += (output.pop(0),)
+            else:
+                attentions = None
+
             return BaseModelOutputWithPooling(
                 last_hidden_state=last_hidden_state,
                 pooler_output=pooler_output,
                 hidden_states=hidden_states,
+                attentions=attentions,
             )

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

@@ -41,6 +41,28 @@ class T5EncoderWrapper(torch.nn.Module):
 
 
 class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
+    """
+    The T5 Model transformer with an encoder-only architecture for feature extraction.
+    This model inherits from [`RBLNTransformerEncoderForFeatureExtraction`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    Important Note:
+        This model supports various sizes of the T5EncoderModel. For optimal performance, it is highly recommended to adjust the tensor parallelism setting
+        based on the model size. Please refer to the [Optimum RBLN Overview](../../../optimum_rbln.md) for guidance on choosing the appropriate tensor parallelism size for your model.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNT5EncoderModel
+
+        model = RBLNT5EncoderModel.from_pretrained(
+            "sentence-transformers/sentence-t5-xxl",
+            export=True,
+            rbln_tensor_parallel_size=4,
+        )
+
+        model.save_pretrained("compiled-sentence-t5-xxl")
+        ```
+    """
+
     auto_model_class = AutoModelForTextEncoding
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
@@ -50,10 +72,7 @@ class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
 
     @classmethod
     def update_rbln_config_using_pipe(
-        cls,
-        pipe: "RBLNDiffusionMixin",
-        rbln_config: "RBLNDiffusionMixinConfig",
-        submodule_name: str,
+        cls, 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
@@ -62,6 +81,29 @@ class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
 
 
 class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
+    """
+    The T5 Model transformer with a language modeling head for conditional generation.
+    This model inherits from [`RBLNModelForSeq2SeqLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    Important Note:
+        This model supports various sizes of the T5ForConditionalGeneration. For optimal performance, it is highly recommended to adjust the tensor parallelism setting
+        based on the model size. Please refer to the [Optimum RBLN Overview](../../../optimum_rbln.md) for guidance on choosing the appropriate tensor parallelism size for your model.
+
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNT5ForConditionalGeneration
+
+        model = RBLNT5ForConditionalGeneration.from_pretrained(
+            "google-t5/t5-11b",
+            export=True,
+            rbln_tensor_parallel_size=4,
+        )
+
+        model.save_pretrained("compiled-sentence-t5-xxl")
+        ```
+    """
+
     support_causal_attn = False
 
     @classmethod

optimum/rbln/transformers/models/{time_series_transformers → time_series_transformer}/__init__.py

@@ -23,4 +23,4 @@
 
 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
+from .modeling_time_series_transformer import RBLNTimeSeriesTransformerForPrediction

optimum/rbln/transformers/models/{time_series_transformers → time_series_transformer}/configuration_time_series_transformer.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -10,7 +10,7 @@ class RBLNTimeSeriesTransformerForPredictionConfig(RBLNModelConfig):
         enc_max_seq_len: Optional[int] = None,
         dec_max_seq_len: Optional[int] = None,
         num_parallel_samples: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

optimum/rbln/transformers/models/{time_series_transformers/modeling_time_series_transformers.py → time_series_transformer/modeling_time_series_transformer.py}

@@ -120,6 +120,17 @@ class RBLNSeq2SeqTSDecoderOutput(ModelOutput):
 
 
 class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
+    """
+    The Time Series Transformer Model with a distribution head on top for time-series forecasting. e.g., for datasets like M4, NN5, or other time series forecasting benchmarks.
+    This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformer-based `TimeSeriesTransformerForPrediction` models on RBLN devices.
+    It implements the methods to convert a pre-trained transformers `TimeSeriesTransformerForPrediction` 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.
+    """
+
     auto_model_class = None
     main_input_name = "inputs_embeds"
 
@@ -144,11 +155,6 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
         )
 
     def __getattr__(self, __name: str) -> Any:
-        """This is the key method to implement RBLN-TimeSeriesTransformersForPrediction.
-        Returns:
-            Any: TimeSeriesTransformersForPrediction's corresponding method
-        """
-
         def redirect(func):
             return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
 
@@ -211,10 +217,9 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
         subfolder: str,
         rbln_config: RBLNTimeSeriesTransformerForPredictionConfig,
     ):
-        """
-        If you are unavoidably running on a CPU rather than an RBLN device,
-        store the torch tensor, weight, etc. in this function.
-        """
+        # If you are unavoidably running on a CPU rather than an RBLN device,
+        # store the torch tensor, weight, etc. in this function.
+
         save_dict = {}
         save_dict["embedder"] = model.model.embedder.state_dict()
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")