optimum-rbln 0.9.4a2__py3-none-any.whl → 0.9.5a4__py3-none-any.whl

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

@@ -14,14 +14,7 @@
 
 from typing import TYPE_CHECKING
 
-import torch.nn as nn
-
-from ...models.decoderonly.decoderonly_architecture import (
-    DecoderOnlyAttention,
-    DecoderOnlyLayer,
-    DecoderOnlyModel,
-    DecoderOnlyWrapper,
-)
+from ...models.decoderonly.decoderonly_architecture import DecoderOnlyWrapper
 
 
 if TYPE_CHECKING:
@@ -31,44 +24,8 @@ if TYPE_CHECKING:
 class OPTWrapper(DecoderOnlyWrapper):
     _use_learned_pos_emb = True
 
-    def get_rbln_attn_class(self):
-        return OPTAttention
-
-    def get_rbln_layer_class(self):
-        return OPTDecoderLayer
-
-    def get_rbln_model_class(self):
-        return OPTModel
-
     def get_model_layer(self, model: "OPTForCausalLM"):
         return model.model.decoder if self.is_causal_lm else model.decoder
 
     def get_decoder_layers(self, model: "OPTForCausalLM"):
         return model.model.decoder.layers if self.is_causal_lm else model.decoder.layers
-
-
-class OPTAttention(DecoderOnlyAttention):
-    def __post_init__(self):
-        self.k_proj = self._original_mod.k_proj
-        self.v_proj = self._original_mod.v_proj
-        self.q_proj = self._original_mod.q_proj
-        self.o_proj = self._original_mod.out_proj
-
-
-class OPTModel(DecoderOnlyModel):
-    def get_embedding(self) -> nn.Embedding:
-        return self._original_mod.embed_tokens
-
-    def get_pos_embedding(self):
-        return self._original_mod.embed_positions
-
-    def get_last_layernorm(self) -> nn.LayerNorm:
-        return self._original_mod.final_layer_norm
-
-
-class OPTDecoderLayer(DecoderOnlyLayer):
-    def get_pre_attention_layernorm(self) -> nn.LayerNorm:
-        return self._original_mod.self_attn_layer_norm
-
-    def get_post_attention_layernorm(self) -> nn.LayerNorm:
-        return self._original_mod.final_layer_norm

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

@@ -0,0 +1,16 @@
+# 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_paligemma import RBLNPaliGemmaForConditionalGenerationConfig, RBLNPaliGemmaModelConfig
+from .modeling_paligemma import RBLNPaliGemmaForConditionalGeneration, RBLNPaliGemmaModel

optimum/rbln/transformers/models/paligemma/configuration_paligemma.py

@@ -0,0 +1,129 @@
+# 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 typing import Any, Optional
+
+from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class RBLNPaliGemmaForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNPaliGemmaForConditionalGenerationConfig.
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized PaliGemma models for multimodal conditional generation tasks
+    that combine vision and language processing capabilities.
+    """
+
+    submodules = ["vision_tower", "language_model"]
+    _allow_no_compile_cfgs = True
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        vision_tower: Optional[RBLNModelConfig] = None,
+        language_model: Optional[RBLNModelConfig] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+                This can include settings specific to the vision encoder, such as input resolution or other vision-related parameters.
+                If not provided, default settings will be used.
+            language_model (Optional[RBLNModelConfig]): Configuration for the language model component.
+                This can include settings specific to the language model, such as tensor parallelism or other text-related parameters.
+                If not provided, default settings will be used.
+            output_hidden_states (Optional[bool]): Whether to output the hidden states of the decoder. Defaults to False.
+            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}")
+
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for PaliGemma vision tower. It will be set to 1.")
+
+        self.output_hidden_states = output_hidden_states or False
+
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower,
+            batch_size=1,  # vision_tower batch_size is always 1 in PaliGemma
+            force_kwargs=True,
+        )
+        self.language_model = self.initialize_submodule_config(
+            submodule_config=language_model,
+            batch_size=batch_size,
+            use_position_ids=True,
+            use_attention_mask=True,
+            use_inputs_embeds=True,
+        )
+
+
+class RBLNPaliGemmaModelConfig(RBLNModelConfig):
+    submodules = ["vision_tower", "language_model"]
+    _allow_no_compile_cfgs = True
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        vision_tower: Optional[RBLNModelConfig] = None,
+        language_model: Optional[RBLNModelConfig] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+                This can include settings specific to the vision encoder, such as input resolution or other vision-related parameters.
+                If not provided, default settings will be used.
+            language_model (Optional[RBLNModelConfig]): Configuration for the language model component.
+                This can include settings specific to the language model, such as tensor parallelism or other text-related parameters.
+                If not provided, default settings will be used.
+            output_hidden_states (Optional[bool]): Whether to output the hidden states of the decoder. Defaults to False.
+            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}")
+
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for PaliGemma vision tower. It will be set to 1.")
+
+        self.output_hidden_states = output_hidden_states or False
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower,
+            batch_size=1,  # vision_tower batch_size is always 1 in PaliGemma
+            force_kwargs=True,
+        )
+
+        self.language_model = self.initialize_submodule_config(
+            submodule_config=language_model,
+            batch_size=batch_size,
+            use_position_ids=True,
+            use_attention_mask=True,
+            use_inputs_embeds=True,
+            output_hidden_states=output_hidden_states,
+        )